Map-making with fixed tones

sotodlib/core/axisman.py

@@ -870,25 +870,78 @@ def reindex_axis(self, axis, indexes, in_place=True):
             else:
                 # By this point we have a non AxisManager
                 # assignment assigned to only our axis.
-                # Build new array with the correct indexes.
-                shape = [len(indexes)]
-                if isinstance(v, np.ndarray):
-                    for s in np.shape(v)[1:]:
-                        shape.append(s)
-
-                new_v = np.empty(shape, dtype=v.dtype)
-                if isinstance(v.dtype, float):
-                    # Fill any float arrays with nans
-                    # Non float arrays may have weird
-                    # behavior for newly added indexes. 
-                    # Oh well.
-                    new_v *= np.nan  
-
-                for i, index in enumerate(indexes):
-                    if np.isnan(index) or not (0 <= index < len(v)):
-                        continue
 
-                    new_v[i] = v[int(index)]
+                from so3g.proj import RangesMatrix, Ranges
+                import scipy.sparse as sp
+
+                if isinstance(v, RangesMatrix):
+                    # RangesMatrix has no dtype; build a new one by selecting
+                    # ranges, using empty Ranges for missing indices.
+                    if len(v.shape) >= 2:
+                        nsamps = v.shape[-1]
+                    elif len(v.ranges) > 0 and hasattr(v.ranges[0], 'count'):
+                        nsamps = v.ranges[0].count
+                    else:
+                        nsamps = 0
+
+                    new_ranges = []
+                    for index in indexes:
+                        if (isinstance(index, float) and np.isnan(index)) \
+                                or not (0 <= int(index) < len(v.ranges)):
+                            new_ranges.append(Ranges(nsamps))
+                        else:
+                            new_ranges.append(v.ranges[int(index)])
+                    new_v = RangesMatrix(new_ranges)
+
+                elif sp.issparse(v):
+                    # Sparse matrix: select rows without densifying.
+                    # Invalid indices become all-zero rows.
+                    v_csr = v.tocsr()
+                    n_rows_in = v_csr.shape[0]
+                    valid_mask = np.array([
+                        not (isinstance(idx, float) and np.isnan(idx))
+                        and 0 <= int(idx) < n_rows_in
+                        for idx in indexes
+                    ])
+                    safe = np.array([
+                        int(idx) if valid else 0
+                        for idx, valid in zip(indexes, valid_mask)
+                    ], dtype=np.int64)
+                    selected = v_csr[safe]
+                    if not valid_mask.all():
+                        diag = sp.diags(valid_mask.astype(v_csr.dtype))
+                        selected = (diag @ selected).tocsr()
+                    new_v = selected
+
+                else:
+                    # ndarray (or anything else with len/shape/dtype/__getitem__)
+                    shape = [len(indexes)]
+                    if isinstance(v, np.ndarray):
+                        for s in np.shape(v)[1:]:
+                            shape.append(s)
+                    elif hasattr(v, 'shape') and len(v.shape) > 1:
+                        for s in v.shape[1:]:
+                            shape.append(s)
+
+                    if v.dtype == float:
+                        # Fill any float arrays with nans
+                        # Non float arrays may have weird
+                        # behavior for newly added indexes.
+                        # Oh well.
+                        new_v = np.full(shape, np.nan)
+                    elif v.dtype == int:
+                        # Fill with -1 values which generally
+                        # correspond ot unassigned (e.g. bias lines)
+                        new_v = np.full(shape, -1)
+                    else:
+                        new_v = np.zeros(shape, dtype=v.dtype)
+
+                    v_len = v.shape[0] if hasattr(v, 'shape') and len(v.shape) > 0 else len(v)
+                    for i, index in enumerate(indexes):
+                        if np.isnan(index) or not (0 <= index < v_len):
+                            continue
+
+                        new_v[i] = v[int(index)]
 
             reindexed_vs[assignment] = new_v
             new_axes[assignment] = np.array(axes)

sotodlib/core/context.py

@@ -356,30 +356,36 @@ def get_obs(self,
 
             # Both tones and det_info exist.
             else:
-                # Grab all band and channel info for dets + tdets
+                # Grab all stream, band, and channel info for dets + tdets
+                det_streams = aman.det_info.stream_id
                 det_bands = aman.det_info.smurf.band
                 det_channels = aman.det_info.smurf.channel
+                tdet_streams = aman.tones.stream_id
                 tdet_bands = aman.tones.band
                 tdet_channels = aman.tones.channel
 
                 # Create a sorted array of dets + tdets
-                special_band_ch = [(b, c) for b, c in zip(tdet_bands, tdet_channels)]
-                normal_band_ch = [(b, c) for b, c in zip(det_bands, det_channels)]
-                band_ch = np.array(sorted(normal_band_ch + special_band_ch))
+                special_band_ch = [(s, b, c) for s, b, c in zip(tdet_streams, tdet_bands, tdet_channels)]
+                normal_band_ch = [(s, b, c) for s, b, c in zip(det_streams, det_bands, det_channels)]
+                band_ch = sorted(normal_band_ch + special_band_ch)
 
                 # Grab the det idxs from the det band + channels
                 det_indexes = np.full(len(band_ch), np.nan)
-                for i, (b, c) in enumerate(band_ch):
-                    w = np.where((det_bands == b) & (det_channels == c))[0]
+                for i, (s, b, c) in enumerate(band_ch):
+                    w = np.where((det_streams == s) & \
+                                 (det_bands == b) & \
+                                 (det_channels == c))[0]
                     if len(w) == 0:
                         continue
 
                     det_indexes[i] = w[0]
 
                 # Grab the tdet idxs from the tdet band + channels
                 tdet_indexes = np.full(len(band_ch), np.nan)
-                for i, (b, c) in enumerate(band_ch):
-                    w = np.where((tdet_bands == b) & (tdet_channels == c))[0]
+                for i, (s, b, c) in enumerate(band_ch):
+                    w = np.where((tdet_streams == s) & \
+                                 (tdet_bands == b) & \
+                                 (tdet_channels == c))[0]
                     if len(w) == 0:
                         continue
 
@@ -393,13 +399,20 @@ def get_obs(self,
 
                 # Finally use the tdet idxs to fill in the tdet data
                 # For the signal, band, and channels
+
+                # If there is no signal, need to pre-populate the signal axis
+                if no_signal:
+                    del aman['signal']
+                    aman.wrap_new('signal', ('dets', 'samps'), dtype='float32')
                 for i, tidx in enumerate(tdet_indexes):
                     if np.isnan(tidx):
                         continue
-    
+
                     aman.signal[i] = aman.tones.signal[int(tidx)]
+                    aman.det_info.stream_id[i] = aman.tones.stream_id[int(tidx)]
                     aman.det_info.smurf.channel[i] = aman.tones.channel[int(tidx)]
                     aman.det_info.smurf.band[i] = aman.tones.band[int(tidx)]
+                    aman.det_info.wafer.type[i] = 'PROB'
 
                 def add_tdet_ids(aman, tdet_indexes, tdet_ids):
                     """

sotodlib/core/metadata/loader.py

@@ -323,7 +323,10 @@ def load_one(self, spec, request, det_info):
         if len(results) == 1:
             result = results[0]
         else:
-            result = results[0].concatenate(results)
+            if isinstance(results[0], core.AxisManager):
+                result = results[0].concatenate(results, other_fields='first')
+            else:
+                result = results[0].concatenate(results)
         return result
 
     def load(self, spec_list, request, det_info=None, free_tags=[],

sotodlib/preprocess/preprocess_util.py

@@ -587,6 +587,7 @@ def load_and_preprocess(obs_id, configs, context=None, dets=None, meta=None,
         pipe = Pipeline(configs["process_pipe"], logger=logger)
         aman = context.get_obs(meta, no_signal=no_signal)
         pipe.run(aman, aman.preprocess, select=False)
+
         return aman, full_aman