Matthews comments, plus go to G3VectorQuat instead of numpy array

docs/proj.rst

@@ -85,9 +85,9 @@ quaternions into rotation angles::
   >>> csl.Q
   spt3g.core.G3VectorQuat([(-0.0384748,0.941783,0.114177,0.313891)])
 
-  >>> csl.coords()   
+  >>> csl.coords()
   array([[ 0.24261138, -0.92726871, -0.99999913, -0.00131952]])
-  
+
 The :func:`coords() <CelestialSightLine.coords>` returns an array with
 shape (n_time, 4); each 4-tuple contains values ``(lon, lat,
 cos(gamma), sin(gamma))``.  The ``lon`` and ``lat`` are the celestial
@@ -117,12 +117,12 @@ orientations::
 
 This particular function, :func:`from_xieta()
 <FocalPlane.from_xieta>`, will apply the SO standard coordinate
-definitions and stores quaternion coefficients (``.quats``) and
-responsivities (``.resps``) for each detectors. These are numpy
-arrays with shape ``[ndet,4]`` and ``[ndet,2]`` respectively.
-``fp.quats[0]`` gives the 4 quaternion coefficients for the
-first detector, while ``fp.resps[0]`` gives its total intensity
-and polarization responsivity.::
+definitions and stores quaternions (``.quats``) and
+responsivities (``.resps``) for each detector. These are a
+``G3VectorQuat``  wth length ``ndet`` and a numpy array
+with shape ``[ndet,2]`` respectively. ``fp.quats[0]`` gives
+the quaternion for the first detector, while ``fp.resps[0]``
+gives its total intensity and polarization responsivity.::
 
   >>> fp.quats[2]
   array([  0.86601716,  0.00377878, -0.00218168,  0.49999524])
@@ -133,7 +133,7 @@ As you can see, the default responsivity is 1 for both total
 intensty and polarization. Note that ``.quats`` contains
 quaternion *coefficients*, not quaternion *objects*. To do
 quaternion math, you need to convert them to actual quaternion
-objects, e.g. ``q = spt3g.core.quat(*fp.quats[0])``, or convert
+objects, e.g. ``q = fp.quats[0]``, or convert
 them all at once with ``qs = spt3g.core.G3VectorQuat(fp.quats)```.
 
 To represent detectors with responsivity different from 1,
@@ -166,7 +166,7 @@ At this point you could get the celestial coordinates for any one of
 those detectors::
 
   # Get vector of quaternion pointings for detector 0
-  q_total = csl.Q * spt3g.core.quat(*fp.quats[0])
+  q_total = csl.Q * fp.quats[0]
   # Decompose q_total into lon, lat, roll angles
   ra, dec, gamma = so3g.proj.quat.decompose_lonlat(q_total)
 
@@ -185,9 +185,8 @@ FocalPlane object as first argument::
    array([[ 0.24261138, -0.92726871,  0.86536489,  0.5011423 ]]),
    array([[ 0.22806774, -0.92722945,  0.50890634,  0.86082189]])]
 
-To be clear, ``coords()`` now returns a dictionary whose keys are the
-detector names.  Each value is an array with shape (n_time,4), and at
-each time step the 4 elements of the array are: ``(lon, lat,
+So ``coords()`` returns a list of numpy arrays with shape (n_time,4),
+and at each time step the 4 elements of the array are: ``(lon, lat,
 cos(gamma), sin(gamma))``.
 
 

python/proj/coords.py

@@ -219,9 +219,9 @@ def coords(self, fplane=None, output=None):
             buffer protocol.
 
         Returns:
-          If fplane is not None, then the result will be
+          If fplane is None, then the result will be
           [n_samp,{lon,lat,cos2psi,sin2psi}]. Otherwise it will
-          be [n_det,n_Samp,{lon,lat,cos2psi,sin2psi}]
+          be [n_det,n_samp,{lon,lat,cos2psi,sin2psi}]
         """
         # Get a projector, in CAR.
         p = so3g.ProjEng_CAR_TQU_NonTiled((1, 1, 1., 1., 1., 1.))
@@ -231,7 +231,7 @@ def coords(self, fplane=None, output=None):
             fplane = FocalPlane.boresight()
             if output is not None:
                 output = output[None]
-        output = p.coords(self.Q, fplane.quats, output)
+        output = p.coords(self.Q, fplane.coeffs(), output)
         if collapse:
             output = output[0]
         return output
@@ -241,11 +241,9 @@ class FocalPlane:
     polarization responses in the focal plane.
 
     Members:
-     quats: Array of float64 with shape [ndet,4] representing the
-      pointing quaternions for each detector. Since these are just
-      plain coefficients, they will need to be converted to quaternion
-      objects if you want to do quaternion math with them, e.g.
-      G3VectorQuat(focal_plane.quats).
+     quats: G3VectorQuat representing the pointing quaternions for
+      each detector. Can be turned into a numpy array of coefficients
+      with np.array(). Or call .coeffs() to get them directly.
      resps: Array of float64 with shape [ndet,2] representing the
       total intensity and polarization responsivity of each detector
      ndet: The number of detectors (read-only)
@@ -274,8 +272,10 @@ def __init__(self, quats=None, resps=None, dets=None):
         # quats will be an quat coeff array-2 and resps will be a numpy
         # array with the right shape, so we don't need to check
         # for this when we use FocalPlane later
-        if quats is None: self.quats = np.zeros((0,4),np.float64)
-        else:             self.quats = np.array(quat.G3VectorQuat(quats))
+        if quats is None: quats = []
+        # Asarray needed because G3VectorQuat doesn't handle list of lists,
+        # which we want to be able to accept
+        self.quats = quat.G3VectorQuat(np.asarray(quats))
         self.resps = np.ones((len(self.quats),2),np.float32)
         if resps is not None:
             self.resps[:] = resps
@@ -286,15 +286,22 @@ def __init__(self, quats=None, resps=None, dets=None):
         # FIXME: old sotodlib compat - remove later
         self._dets   = list(dets) if dets is not None else []
         self._detmap = {name:i for i,name in enumerate(self._dets)}
+    def coeffs(self):
+        """Return an [ndet,4] numpy array representing the quaternion
+        coefficients of ``.quats``. Useful for passing the detector
+        pointing to C++ code"""
+        return np.array(self.quats, dtype=np.float64)
     @classmethod
     def boresight(cls):
         """Construct a FocalPlane representing a single detector with a
         responsivity of one pointing along the telescope boresight"""
-        quats = np.array([[1,0,0,0]],np.float64)
-        return cls(quats=quats)
+        return cls(quats=[[1,0,0,0]])
+    # FIXME: old sotodlib compat - expand to actual argument list later
     @classmethod
-    def from_xieta(cls, xi, eta, gamma=0, T=1, P=1, Q=1, U=0, hwp=False):
-        """Construct a FocalPlane from focal plane coordinates (xi,eta).
+    def from_xieta(cls, *args, **kwargs):
+        """
+        def from_xieta(cls, xi, eta, gamma=0, T=1, P=1, Q=1, U=0, hwp=False):
+        Construct a FocalPlane from focal plane coordinates (xi,eta).
 
         These are Cartesian projection plane coordinates. When looking
         at the sky along the telescope boresight, xi is parallel to
@@ -346,7 +353,7 @@ def from_xieta(cls, xi, eta, gamma=0, T=1, P=1, Q=1, U=0, hwp=False):
         # response too. Writing it like this handles both cases, and
         # as a bonus(?) any combination of them
         # FIXME: old sotodlib compat - remove later
-        xi, eta, gamma, T, P, Q, U, hwp, dets = cls._xieta_compat(xi,eta,gamma,T,P,Q,U,hwp)
+        xi, eta, gamma, T, P, Q, U, hwp, dets = cls._xieta_compat(*args, **kwargs)
         gamma = gamma + np.arctan2(U,Q)/2
         P     = P * (Q**2+U**2)**0.5
         if hwp: gamma = -gamma
@@ -359,53 +366,57 @@ def from_xieta(cls, xi, eta, gamma=0, T=1, P=1, Q=1, U=0, hwp=False):
         # FIXME: old sotodlib compat - remove dets argument later
         return cls(quats, resps=resps, dets=dets)
     def __repr__(self):
-        return "FocalPlane(quats=%s,resps=%s)" % (repr(self.quats), repr(self.resps))
+        return "FocalPlane(quats=%s,resps=%s)" % (repr(self.coeffs()), repr(self.resps))
     @property
     def ndet(self): return len(self.quats)
     def __len__(self): return self.ndet
     def __getitem__(self, sel):
         """Slice the FocalPlane with slice sel, resulting in a new
-        FocalPlane with a subset of the detectors. Only 1d slice supported,
-        not integers or multidimensional slices. Example: ``focal_plane[10:20]``
-        would make a sub-FocalPlane with detector indices 10,11,...,19.
+        FocalPlane with a subset of the detectors. Only a 1d slice
+        or boolean mask is supported, not integers or multidimensional
+        slices. Example: ``focal_plane[10:20]`` would make a
+        sub-FocalPlane with detector indices 10,11,...,19.
 
         Deprecated: Temporarily also supports that sel is a detector name,
         in which case an  ``spt3g.core.quat`` is returned for that detector.
         This is provided for backwards compatibility."""
         # FIXME: old sotodlib compat - remove later
-        if isinstance(sel, str): return quat.G3VectorQuat(self.quats)[self._dets.index(sel)]
-        return FocalPlane(quats=self.quats[sel], resps=self.resps[sel])
+        if isinstance(sel, str): return self.quats[self._dets.index(sel)]
+        # We go via .coeffs() here to get around G3VectorQuat's lack
+        # of boolean mask support
+        return FocalPlane(quats=self.coeffs()[sel], resps=self.resps[sel])
     def items(self):
         """Iterate over detector quaternions and responsivities. Yields
         ``(spt3g.core.quat, array([Tresp,Presp]))`` pairs. Unlke the raw
         entries in the .quats member, which are just numpy arrays with
         length 4,  ``spt3g.core.quat`` are proper quaternon objects that
         support quaternion math.
         """
-        for q, resp in zip(quat.G3VectorQuat(self.quats), self.resps):
+        for q, resp in zip(self.quats, self.resps):
             yield q, resp
     # FIXME: old sotodlib compat - remove later
-    @classmethod
-    def _xieta_compat(cls, xi, eta, gamma, T, P, Q, U, hwp):
+    @staticmethod
+    def _xieta_compat(*args, **kwargs):
         # Accept the alternative format (names,xi,eta,gamma)
-        if isinstance(xi[0], str):
-            return eta, gamma, T, 1, 1, 1, 0, False, xi
+        def helper(xi, eta, gamma=0, T=1, P=1, Q=1, U=0, hwp=False, dets=None):
+            return xi, eta, gamma, T, P, Q, U, hwp, dets
+        if isinstance(args[0][0], str):
+            return helper(*args[1:], dets=args[0], **kwargs)
         else:
-            return xi, eta, gamma, T, P, Q, U, hwp, None
+            return helper(*args, **kwargs)
     # FIXME: old sotodlib compat - remove later
     def __setitem__(self, name, q):
         # Make coords/pmat.py _get_asm work in old sotodlib. It
         # expects to be able to build up a focalplane by assigning
         # quats one at a time
-        q = np.array(quat.G3VectorQuat([q]))
         if name in self._detmap:
             i = self._detmap[i]
-            self.quats[i] = q[0]
+            self.quats[i] = q
         else:
             self._dets.append(name)
             self._detmap[name] = len(self._dets)-1
+            self.quats.append(q)
             # This is inefficient, but it's temporary
-            self.quats = np.concatenate([self.quats,q])
             self.resps = np.concatenate([self.resps,np.ones((1,2),np.float32)])
 
 class Assembly:

python/proj/wcs.py

@@ -183,7 +183,7 @@ def get_pixels(self, assembly):
         """
         projeng = self.get_ProjEng('TQU')
         q_native = self._cache_q_fp_to_native(assembly.Q)
-        return projeng.pixels(q_native, assembly.fplane.quats, None)
+        return projeng.pixels(q_native, assembly.fplane.coeffs(), None)
 
     def get_pointing_matrix(self, assembly):
         """Get the pointing matrix information, which is to say both the pixel
@@ -199,7 +199,7 @@ def get_pointing_matrix(self, assembly):
         """
         projeng = self.get_ProjEng('TQU')
         q_native = self._cache_q_fp_to_native(assembly.Q)
-        return projeng.pointing_matrix(q_native, assembly.fplane.quats, assembly.fplane.resps, None, None)
+        return projeng.pointing_matrix(q_native, assembly.fplane.coeffs(), assembly.fplane.resps, None, None)
 
     def get_coords(self, assembly, use_native=False, output=None):
         """Get the spherical coordinates for the provided pointing Assembly.
@@ -225,7 +225,7 @@ def get_coords(self, assembly, use_native=False, output=None):
             q_native = self._cache_q_fp_to_native(assembly.Q)
         else:
             q_native = assembly.Q
-        return projeng.coords(q_native, assembly.fplane.quats, output)
+        return projeng.coords(q_native, assembly.fplane.coeffs(), output)
 
     def get_planar(self, assembly, output=None):
         """Get projection plane coordinates for all detectors at all times.
@@ -240,7 +240,7 @@ def get_planar(self, assembly, output=None):
         """
         q_native = self._cache_q_fp_to_native(assembly.Q)
         projeng = self.get_ProjEng('TQU')
-        return projeng.coords(q_native, assembly.fplane.quats, output)
+        return projeng.coords(q_native, assembly.fplane.coeffs(), output)
 
     def zeros(self, super_shape):
         """Return a map, filled with zeros, with shape (super_shape,) +
@@ -272,7 +272,7 @@ def to_map(self, signal, assembly, output=None, det_weights=None,
             comps = self._guess_comps(self._get_map_shape(output))
         projeng = self.get_ProjEng(comps)
         q_native = self._cache_q_fp_to_native(assembly.Q)
-        map_out = projeng.to_map(output, q_native, assembly.fplane.quats,
+        map_out = projeng.to_map(output, q_native, assembly.fplane.coeffs(),
             assembly.fplane.resps, signal, det_weights, threads)
         return map_out
 
@@ -298,7 +298,7 @@ def to_weights(self, assembly, output=None, det_weights=None,
             comps = self._guess_comps(shape[1:])
         projeng = self.get_ProjEng(comps)
         q_native = self._cache_q_fp_to_native(assembly.Q)
-        map_out = projeng.to_weight_map(output, q_native, assembly.fplane.quats,
+        map_out = projeng.to_weight_map(output, q_native, assembly.fplane.coeffs(),
             assembly.fplane.resps, det_weights, threads)
         return map_out
 
@@ -318,7 +318,7 @@ def from_map(self, src_map, assembly, signal=None, comps=None):
             comps = self._guess_comps(self._get_map_shape(src_map))
         projeng = self.get_ProjEng(comps)
         q_native = self._cache_q_fp_to_native(assembly.Q)
-        signal_out = projeng.from_map(src_map, q_native, assembly.fplane.quats,
+        signal_out = projeng.from_map(src_map, q_native, assembly.fplane.coeffs(),
             assembly.fplane.resps, signal)
         return signal_out
 
@@ -362,7 +362,7 @@ def assign_threads(self, assembly, method='domdir', n_threads=None):
         if method == 'simple':
             projeng = self.get_ProjEng('T')
             q_native = self._cache_q_fp_to_native(assembly.Q)
-            omp_ivals = projeng.pixel_ranges(q_native, assembly.fplane.quats, None, n_threads)
+            omp_ivals = projeng.pixel_ranges(q_native, assembly.fplane.coeffs(), None, n_threads)
             return wrap_ivals(omp_ivals)
 
         elif method == 'domdir':
@@ -402,7 +402,7 @@ def assign_threads_from_map(self, assembly, tmap, n_threads=None):
         projeng = self.get_ProjEng('T')
         q_native = self._cache_q_fp_to_native(assembly.Q)
         n_threads = mapthreads.get_num_threads(n_threads)
-        omp_ivals = projeng.pixel_ranges(q_native, assembly.fplane.quats, tmap, n_threads)
+        omp_ivals = projeng.pixel_ranges(q_native, assembly.fplane.coeffs(), tmap, n_threads)
         return wrap_ivals(omp_ivals)
 
     def get_active_tiles(self, assembly, assign=False):
@@ -441,7 +441,7 @@ def get_active_tiles(self, assembly, assign=False):
         projeng = self.get_ProjEng('T')
         q_native = self._cache_q_fp_to_native(assembly.Q)
         # This returns a G3VectorInt of length n_tiles giving count of hits per tile.
-        hits = np.array(projeng.tile_hits(q_native, assembly.fplane.quats))
+        hits = np.array(projeng.tile_hits(q_native, assembly.fplane.coeffs()))
         tiles = np.nonzero(hits)[0]
         hits = hits[tiles]
         if assign is True:
@@ -460,7 +460,7 @@ def get_active_tiles(self, assembly, assign=False):
             #     print(f"Warning: Threads poorly balanced. Max/mean hits = {max_ratio}")
 
             # Now paint them into Ranges.
-            R = projeng.tile_ranges(q_native, assembly.fplane.quats, group_tiles)
+            R = projeng.tile_ranges(q_native, assembly.fplane.coeffs(), group_tiles)
             R = wrap_ivals(R)
             return {
                 'active_tiles': list(tiles),
@@ -795,7 +795,7 @@ def get_coords(self, assembly, use_native=False, output=None):
             q_native = self._cache_q_fp_to_native(assembly.Q)
         else:
             q_native = assembly.Q
-        return projeng.coords(q_native, assembly.dets, output)
+        return projeng.coords(q_native, assembly.fplane.coeffs(), output)
 
     def from_map(self, src_map, assembly, signal=None, comps=None):
         """De-project from a map, returning a Signal-like object.

test/test_proj_eng.py

@@ -73,7 +73,9 @@ def test_00_basic(self):
                      det_weights=np.array([0., 0.], dtype='float32'))[0]
         assert(np.all(m==0))
 
-        asm.fplane.quats[1,2] = np.nan
+        # Can't assign to quat fields, so do
+        # it this way instead
+        asm.fplane.quats[1] = asm.fplane.quats[1]*np.nan
         with self.assertRaises(ValueError):
            p.to_map(sig, asm, comps='T')
         with self.assertRaises(ValueError):

test/test_proj_eng_hp.py

@@ -32,7 +32,7 @@ class TestProjEngHP(unittest.TestCase):
     """Test ProjectionistHealpix
        Based on TestProjEng
     """
-    
+
     def test_00_basic(self):
         scan, asm = get_basics()
         nside = 128
@@ -51,8 +51,9 @@ def test_00_basic(self):
                      det_weights=np.array([0., 0.], dtype='float32'))[0]
         assert(np.all(m==0))
 
-        # Raise if pointing invalid.
-        asm.fplane.quats[1,2] = np.nan
+        # Can't assign to quat fields, so do
+        # it this way instead
+        asm.fplane.quats[1] = asm.fplane.quats[1]*np.nan
         with self.assertRaises(ValueError):
            p.to_map(sig, asm, comps='T')
         with self.assertRaises(ValueError):
@@ -74,7 +75,7 @@ def test_10_tiled(self):
                 assert(np.any(w2))
             # Identify active subtiles?
             print(p.active_tiles)
-        
+
     def test_20_threads(self):
         for (tiled, interpol, method) in itertools.product(
                 [False, True],
@@ -88,7 +89,7 @@ def test_20_threads(self):
                 nside_tile = 8
             else:
                 nside_tile = None
-                
+
             p = proj.ProjectionistHealpix.for_healpix(nside, nside_tile, interpol=interpol)
             sig = np.ones((2, len(scan[0])), 'float32')
             n_threads = 3