From 69a337db28b44d9a52c669f24524e277e5f9656b Mon Sep 17 00:00:00 2001 From: Ioannis Paraskevakos Date: Tue, 20 Aug 2024 13:55:27 -0400 Subject: [PATCH 1/8] linting to debug eventually --- sotodlib/core/context.py | 59 +++++++++++++++++++++++----------------- 1 file changed, 34 insertions(+), 25 deletions(-) diff --git a/sotodlib/core/context.py b/sotodlib/core/context.py index 1cd291027..aeb050ae8 100644 --- a/sotodlib/core/context.py +++ b/sotodlib/core/context.py @@ -5,12 +5,15 @@ import logging import numpy as np +from typing import Union, Dict, Tuple, List + from . import metadata from .util import tag_substr from .axisman import AxisManager, OffsetAxis, AxisInterface logger = logging.getLogger(__name__) + class Context(odict): # Sets of special handlers may be registered in this class variable, then # requested by name in the context.yaml key "context_hooks". @@ -149,19 +152,22 @@ def reload(self, load_list='all'): self.loader \ = metadata.SuperLoader(self, obsdb=self.obsdb) - def get_obs(self, - obs_id=None, - dets=None, - samples=None, - filename=None, - detsets=None, - meta=None, - ignore_missing=None, - on_missing=None, - free_tags=None, - no_signal=None, - loader_type=None, - ): + def get_obs( + self, + obs_id: Union[str, Dict[str, str], AxisManager] = None, + dets: Union[ + List[str], Dict[str, Union[List[str], str]], metadata.ResultSet + ] = None, + samples: Tuple[int] = None, + filename: str = None, + detsets: List[str] = None, + meta: AxisManager = None, + ignore_missing: bool = None, + on_missing: Dict[str, str] = None, + free_tags: List[str] = None, + no_signal: bool = None, + loader_type: str = None, + ) -> AxisManager: """Load TOD and supporting metadata for some observation. Most arguments to this function are also accepted by (and in @@ -311,18 +317,21 @@ def get_obs(self, aman.merge(meta) return aman - def get_meta(self, - obs_id=None, - dets=None, - samples=None, - filename=None, - detsets=None, - meta=None, - free_tags=None, - check=False, - ignore_missing=False, - on_missing=None, - det_info_scan=False): + def get_meta( + self, + obs_id: Union[str, Dict[str, str], AxisManager] = None, + dets: Union[List[str], Dict[str, Union[List[str], str]], ResultSet] = None, + samples: Tuple[int] = None, + filename: str = None, + detsets: List[str] = None, + meta: AxisManager = None, + free_tags: List[str] = None, + check: bool = False, + ignore_missing: bool = False, + on_missing: Dict[str, str] = None, + det_info_scan: bool = False, + ) -> AxisManager: + """Load supporting metadata for an observation and return it in an AxisManager. From e2f259b67a825481c2dc5df7b2bfeb4c6ee49a78 Mon Sep 17 00:00:00 2001 From: chervias Date: Thu, 5 Sep 2024 17:37:26 -0700 Subject: [PATCH 2/8] Fix for flags and downsampling. This will allow the ML/depth1 mapmakers to make maps with downsampling --- sotodlib/mapmaking/utilities.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/sotodlib/mapmaking/utilities.py b/sotodlib/mapmaking/utilities.py index dc65fbbb3..51bdad5cd 100644 --- a/sotodlib/mapmaking/utilities.py +++ b/sotodlib/mapmaking/utilities.py @@ -510,15 +510,17 @@ def downsample_obs(obs, down): res.wrap("signal", resample.resample_fft_simple(obs.signal, onsamp), [(0,"dets"),(1,"samps")]) # The cuts - # TODO: The TOD will include a Flagmanager with all the flags. Update this part - # accordingly. - cut_keys = ["glitch_flag"] + # obs.flags will contain all types of flags. We should query it for glitch_flags and source_flags + cut_keys = ["glitch_flags"] - if "source_flags" in obs: + if "source_flags" in obs.flags: cut_keys.append("source_flags") + # We need to add a res.flags FlagManager to res + res = res.wrap('flags', core.FlagManager.for_tod(res)) + for key in cut_keys: - res.wrap(key, downsample_cut(getattr(obs, key), down), [(0,"dets"),(1,"samps")]) + res.flags.wrap(key, downsample_cut(getattr(obs.flags, key), down), [(0,"dets"),(1,"samps")]) # Not sure how to deal with flags. Some sort of or-binning operation? But it # doesn't matter anyway From 7b3c00b79d6ca62436ad875cba884c013305e4b6 Mon Sep 17 00:00:00 2001 From: Ioannis Paraskevakos Date: Mon, 9 Sep 2024 11:20:53 -0400 Subject: [PATCH 3/8] setting flags input --- sotodlib/mapmaking/ml_mapmaker.py | 6 +++--- sotodlib/mapmaking/utilities.py | 3 +-- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/sotodlib/mapmaking/ml_mapmaker.py b/sotodlib/mapmaking/ml_mapmaker.py index 198198318..917b9d240 100644 --- a/sotodlib/mapmaking/ml_mapmaker.py +++ b/sotodlib/mapmaking/ml_mapmaker.py @@ -202,7 +202,7 @@ def __init__(self, shape, wcs, comm, comps="TQU", name="sky", ofmt="{name}", out self.div = enmap.zeros((ncomp,ncomp)+shape, wcs, dtype=dtype) self.hits= enmap.zeros( shape, wcs, dtype=dtype) - def add_obs(self, id, obs, nmat, Nd, pmap=None): + def add_obs(self, id, obs, nmat, Nd, pmap=None, glitch_flags="flags.glitch_flags"): """Add and process an observation, building the pointing matrix and our part of the RHS. "obs" should be an Observation axis manager, nmat a noise model, representing the inverse noise covariance matrix, @@ -372,7 +372,7 @@ def __init__(self, comm, name="cut", ofmt="{name}_{rank:02}", dtype=np.float32, self.rhs = [] self.div = [] - def add_obs(self, id, obs, nmat, Nd): + def add_obs(self, id, obs, nmat, Nd, glitch_flags="flags.glitch_flags"): """Add and process an observation. "obs" should be an Observation axis manager, nmat a noise model, representing the inverse noise covariance matrix, and Nd the result of applying the noise model to the detector time-ordered data.""" @@ -441,7 +441,7 @@ def translate(self, other, junk): so3g.translate_cuts(odata.pcut.cuts, sdata.pcut.cuts, sdata.pcut.model, sdata.pcut.params, junk[odata.i1:odata.i2], res[sdata.i1:sdata.i2]) return res - def transeval(self, id, obs, other, junk, tod): + def transeval(self, id, obs, other, junk, tod, glitch_flags="flags.glitch_flags"): """Translate data junk from SignalCut other to the current SignalCut, and then evaluate it for the given observation, returning a tod. This is used when building a signal-free tod for the noise model diff --git a/sotodlib/mapmaking/utilities.py b/sotodlib/mapmaking/utilities.py index dc65fbbb3..a099ca454 100644 --- a/sotodlib/mapmaking/utilities.py +++ b/sotodlib/mapmaking/utilities.py @@ -136,7 +136,6 @@ def safe_invert_div(div, lim=1e-2, lim0=np.finfo(np.float32).tiny**0.5): return idiv - def measure_cov(d, nmax=10000): d = d[:,::max(1,d.shape[1]//nmax)] n,m = d.shape @@ -512,7 +511,7 @@ def downsample_obs(obs, down): # The cuts # TODO: The TOD will include a Flagmanager with all the flags. Update this part # accordingly. - cut_keys = ["glitch_flag"] + cut_keys = ["glitch_flags"] if "source_flags" in obs: cut_keys.append("source_flags") From bd0468d9727514cc1fa01f466d559b9cd6567c6b Mon Sep 17 00:00:00 2001 From: Ioannis Paraskevakos Date: Mon, 9 Sep 2024 15:38:11 -0400 Subject: [PATCH 4/8] reverting linting change --- sotodlib/core/context.py | 57 ++++++++++++++++++---------------------- 1 file changed, 26 insertions(+), 31 deletions(-) diff --git a/sotodlib/core/context.py b/sotodlib/core/context.py index 191d87ed1..430d19920 100644 --- a/sotodlib/core/context.py +++ b/sotodlib/core/context.py @@ -152,22 +152,19 @@ def reload(self, load_list='all'): self.loader \ = metadata.SuperLoader(self, obsdb=self.obsdb) - def get_obs( - self, - obs_id: Union[str, Dict[str, str], AxisManager] = None, - dets: Union[ - List[str], Dict[str, Union[List[str], str]], metadata.ResultSet - ] = None, - samples: Tuple[int] = None, - filename: str = None, - detsets: List[str] = None, - meta: AxisManager = None, - ignore_missing: bool = None, - on_missing: Dict[str, str] = None, - free_tags: List[str] = None, - no_signal: bool = None, - loader_type: str = None, - ) -> AxisManager: + def get_obs(self, + obs_id=None, + dets=None, + samples=None, + filename=None, + detsets=None, + meta=None, + ignore_missing=None, + on_missing=None, + free_tags=None, + no_signal=None, + loader_type=None, + ): """Load TOD and supporting metadata for some observation. Most arguments to this function are also accepted by (and in @@ -317,21 +314,19 @@ def get_obs( aman.merge(meta) return aman - def get_meta( - self, - obs_id: Union[str, Dict[str, str], AxisManager] = None, - dets: Union[List[str], Dict[str, Union[List[str], str]], metadata.ResultSet] = None, - samples: Tuple[int] = None, - filename: str = None, - detsets: List[str] = None, - meta: AxisManager = None, - free_tags: List[str] = None, - check: bool = False, - ignore_missing: bool = False, - on_missing: Dict[str, str] = None, - det_info_scan: bool = False, - ) -> AxisManager: - + def get_meta(self, + obs_id=None, + dets=None, + samples=None, + filename=None, + detsets=None, + meta=None, + free_tags=None, + check=False, + ignore_missing=False, + on_missing=None, + det_info_scan=False + ): """Load supporting metadata for an observation and return it in an AxisManager. From 7c44e535e4b2b99b6322f858a41a189d1ced887d Mon Sep 17 00:00:00 2001 From: Ioannis Paraskevakos Date: Mon, 9 Sep 2024 15:56:14 -0400 Subject: [PATCH 5/8] adding missing call --- sotodlib/mapmaking/ml_mapmaker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sotodlib/mapmaking/ml_mapmaker.py b/sotodlib/mapmaking/ml_mapmaker.py index ba221cc5d..c1b12d913 100644 --- a/sotodlib/mapmaking/ml_mapmaker.py +++ b/sotodlib/mapmaking/ml_mapmaker.py @@ -378,7 +378,7 @@ def add_obs(self, id, obs, nmat, Nd, glitch_flags="flags.glitch_flags"): nmat a noise model, representing the inverse noise covariance matrix, and Nd the result of applying the noise model to the detector time-ordered data.""" Nd = Nd.copy() # This copy can be avoided if build_obs is split into two parts - pcut = PmatCut(obs.flags.glitch_flags, model=self.cut_type) + pcut = PmatCut(get_flags_from_path(obs, glitch_flags), model=self.cut_type) # Build our RHS obs_rhs = np.zeros(pcut.njunk, self.dtype) pcut.backward(Nd, obs_rhs) From 51b0e3b89abaaea644e99f0a03708902532ab16a Mon Sep 17 00:00:00 2001 From: Ioannis Paraskevakos Date: Mon, 23 Sep 2024 15:41:55 -0400 Subject: [PATCH 6/8] moving glitch flags path to init --- sotodlib/mapmaking/ml_mapmaker.py | 35 ++++++++++++++++++------------- 1 file changed, 21 insertions(+), 14 deletions(-) diff --git a/sotodlib/mapmaking/ml_mapmaker.py b/sotodlib/mapmaking/ml_mapmaker.py index c1b12d913..2de7f3fe5 100644 --- a/sotodlib/mapmaking/ml_mapmaker.py +++ b/sotodlib/mapmaking/ml_mapmaker.py @@ -1,9 +1,14 @@ import numpy as np -from pixell import bunch, enmap, tilemap, utils +import h5py +import so3g +from pixell import bunch, enmap, tilemap, putils from .. import coords -from .pointing_matrix import * -from .utilities import * +from .pointing_matrix import PmatCut +from .utilities import (MultiZipper, get_flags_from_path, recentering_to_quat_lonlat, + evaluate_recentering, TileMapZipper, MapZipper, + safe_invert_div, unarr, ArrayZipper) +from .noise_model import NmatUncorr class MLMapmaker: @@ -147,7 +152,7 @@ def transeval(self, id, obs, other, x, tod=None): class Signal: """This class represents a thing we want to solve for, e.g. the sky, ground, cut samples, etc.""" - def __init__(self, name, ofmt, output, ext): + def __init__(self, name, ofmt, output, ext, glitch_flags: str = "flags.glitch_flags"): """Initialize a Signal. It probably doesn't make sense to construct a generic signal directly, though. Use one of the subclasses. Arguments: @@ -162,6 +167,7 @@ def __init__(self, name, ofmt, output, ext): self.ext = ext self.dof = None self.ready = False + self.glitch_flags = glitch_flags def add_obs(self, id, obs, nmat, Nd): pass def prepare(self): self.ready = True def forward (self, id, tod, x): pass @@ -177,12 +183,12 @@ class SignalMap(Signal): """Signal describing a non-distributed sky map.""" def __init__(self, shape, wcs, comm, comps="TQU", name="sky", ofmt="{name}", output=True, ext="fits", dtype=np.float32, sys=None, recenter=None, tile_shape=(500,500), tiled=False, - interpol=None): + interpol=None, glitch_flags: str = "flags.glitch_flags"): """Signal describing a sky map in the coordinate system given by "sys", which defaults to equatorial coordinates. If tiled==True, then this will be a distributed map with the given tile_shape, otherwise it will be a plain enmap. interpol controls the pointing matrix interpolation mode. See so3g's Projectionist docstring for details.""" - Signal.__init__(self, name, ofmt, output, ext) + Signal.__init__(self, name, ofmt, output, ext, glitch_flags) self.comm = comm self.comps = comps self.sys = sys @@ -203,7 +209,7 @@ def __init__(self, shape, wcs, comm, comps="TQU", name="sky", ofmt="{name}", out self.div = enmap.zeros((ncomp,ncomp)+shape, wcs, dtype=dtype) self.hits= enmap.zeros( shape, wcs, dtype=dtype) - def add_obs(self, id, obs, nmat, Nd, pmap=None, glitch_flags:str ="flags.glitch_flags"): + def add_obs(self, id, obs, nmat, Nd, pmap=None): """Add and process an observation, building the pointing matrix and our part of the RHS. "obs" should be an Observation axis manager, nmat a noise model, representing the inverse noise covariance matrix, @@ -211,7 +217,7 @@ def add_obs(self, id, obs, nmat, Nd, pmap=None, glitch_flags:str ="flags.glitch_ """ Nd = Nd.copy() # This copy can be avoided if build_obs is split into two parts ctime = obs.timestamps - pcut = PmatCut(get_flags_from_path(obs, glitch_flags)) # could pass this in, but fast to construct + pcut = PmatCut(get_flags_from_path(obs, self.glitch_flags)) # could pass this in, but fast to construct if pmap is None: # Build the local geometry and pointing matrix for this observation if self.recenter: @@ -348,6 +354,7 @@ def transeval(self, id, obs, other, map, tod): # Currently we don't support any actual translation, but could handle # resolution changes in the future (probably not useful though) self._checkcompat(other) + ctime = obs.timestamp # Build the local geometry and pointing matrix for this observation if self.recenter: rot = recentering_to_quat_lonlat(*evaluate_recentering(self.recenter, @@ -362,9 +369,9 @@ def transeval(self, id, obs, other, map, tod): class SignalCut(Signal): def __init__(self, comm, name="cut", ofmt="{name}_{rank:02}", dtype=np.float32, - output=False, cut_type=None): + output=False, cut_type=None, glitch_flags:str ="flags.glitch_flags"): """Signal for handling the ML solution for the values of the cut samples.""" - Signal.__init__(self, name, ofmt, output, ext="hdf") + Signal.__init__(self, name, ofmt, output, ext="hdf", glitch_flags=glitch_flags) self.comm = comm self.data = {} self.dtype = dtype @@ -373,12 +380,12 @@ def __init__(self, comm, name="cut", ofmt="{name}_{rank:02}", dtype=np.float32, self.rhs = [] self.div = [] - def add_obs(self, id, obs, nmat, Nd, glitch_flags="flags.glitch_flags"): + def add_obs(self, id, obs, nmat, Nd): """Add and process an observation. "obs" should be an Observation axis manager, nmat a noise model, representing the inverse noise covariance matrix, and Nd the result of applying the noise model to the detector time-ordered data.""" Nd = Nd.copy() # This copy can be avoided if build_obs is split into two parts - pcut = PmatCut(get_flags_from_path(obs, glitch_flags), model=self.cut_type) + pcut = PmatCut(get_flags_from_path(obs, self.glitch_flags), model=self.cut_type) # Build our RHS obs_rhs = np.zeros(pcut.njunk, self.dtype) pcut.backward(Nd, obs_rhs) @@ -442,7 +449,7 @@ def translate(self, other, junk): so3g.translate_cuts(odata.pcut.cuts, sdata.pcut.cuts, sdata.pcut.model, sdata.pcut.params, junk[odata.i1:odata.i2], res[sdata.i1:sdata.i2]) return res - def transeval(self, id, obs, other, junk, tod, glitch_flags: str="flags.glitch_flags"): + def transeval(self, id, obs, other, junk, tod): """Translate data junk from SignalCut other to the current SignalCut, and then evaluate it for the given observation, returning a tod. This is used when building a signal-free tod for the noise model @@ -450,7 +457,7 @@ def transeval(self, id, obs, other, junk, tod, glitch_flags: str="flags.glitch_f self._checkcompat(other) # We have to make a pointing matrix from scratch because add_obs # won't have been called yet at this point - spcut = PmatCut(get_flags_from_path(obs, glitch_flags), model=self.cut_type) + spcut = PmatCut(get_flags_from_path(obs, self.glitch_flags), model=self.cut_type) # We do have one for other though, since that will be the output # from the previous round of multiplass mapmaking. odata = other.data[id] From 815dabc01fe5417fed3378b8904a4912feb0e317 Mon Sep 17 00:00:00 2001 From: Ioannis Paraskevakos Date: Tue, 24 Sep 2024 09:53:20 -0400 Subject: [PATCH 7/8] fix:syntax error --- sotodlib/mapmaking/ml_mapmaker.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/sotodlib/mapmaking/ml_mapmaker.py b/sotodlib/mapmaking/ml_mapmaker.py index 2de7f3fe5..567470f04 100644 --- a/sotodlib/mapmaking/ml_mapmaker.py +++ b/sotodlib/mapmaking/ml_mapmaker.py @@ -1,7 +1,8 @@ import numpy as np import h5py import so3g -from pixell import bunch, enmap, tilemap, putils +from pixell import bunch, enmap, tilemap +from pixell import utils as putils from .. import coords from .pointing_matrix import PmatCut @@ -42,7 +43,7 @@ def add_obs(self, id, obs, deslope=True, noise_model=None, signal_estimate=None) # the noise model, if available if signal_estimate is not None: tod -= signal_estimate if deslope: - utils.deslope(tod, w=5, inplace=True) + putils.deslope(tod, w=5, inplace=True) # Allow the user to override the noise model on a per-obs level if noise_model is None: noise_model = self.noise_model # Build the noise model from the obs unless a fully @@ -61,7 +62,7 @@ def add_obs(self, id, obs, deslope=True, noise_model=None, signal_estimate=None) # The signal estimate might not be desloped, so # adding it back can reintroduce a slope. Fix that here. if deslope: - utils.deslope(tod, w=5, inplace=True) + putils.deslope(tod, w=5, inplace=True) # And apply it to the tod tod = nmat.apply(tod) # Add the observation to each of our signals @@ -125,7 +126,7 @@ def solve(self, maxiter=500, maxerr=1e-6, x0=None): self.prepare() rhs = self.dof.zip(*[signal.rhs for signal in self.signals]) if x0 is not None: x0 = self.dof.zip(*x0) - solver = utils.CG(self.A, rhs, M=self.M, dot=self.dof.dot, x0=x0) + solver = putils.CG(self.A, rhs, M=self.M, dot=self.dof.dot, x0=x0) while solver.i < maxiter and solver.err > maxerr: solver.step() yield bunch.Bunch(i=solver.i, err=solver.err, x=self.dof.unzip(solver.x)) @@ -268,9 +269,9 @@ def prepare(self): self.dof = TileMapZipper(self.rhs.geometry, dtype=self.dtype, comm=self.comm) else: if self.comm is not None: - self.rhs = utils.allreduce(self.rhs, self.comm) - self.div = utils.allreduce(self.div, self.comm) - self.hits = utils.allreduce(self.hits, self.comm) + self.rhs = putils.allreduce(self.rhs, self.comm) + self.div = putils.allreduce(self.div, self.comm) + self.hits = putils.allreduce(self.hits, self.comm) self.dof = MapZipper(*self.rhs.geometry, dtype=self.dtype) self.idiv = safe_invert_div(self.div) self.ready = True @@ -307,7 +308,7 @@ def from_work(self, map): return tilemap.redistribute(map, self.comm, self.rhs.geometry.active) else: if self.comm is None: return map - else: return utils.allreduce(map, self.comm) + else: return putils.allreduce(map, self.comm) def write(self, prefix, tag, m): if not self.output: return From 0a48177cf04389454cb808e2aafca0625bf96fc0 Mon Sep 17 00:00:00 2001 From: Ioannis Paraskevakos Date: Tue, 24 Sep 2024 15:39:43 -0400 Subject: [PATCH 8/8] glitch flags in MLmapmaker plus some utils changes --- sotodlib/core/g3_core.py | 2 +- sotodlib/mapmaking/ml_mapmaker.py | 24 ++++++---- sotodlib/mapmaking/utilities.py | 76 +++++++++++++++++++++---------- 3 files changed, 69 insertions(+), 33 deletions(-) diff --git a/sotodlib/core/g3_core.py b/sotodlib/core/g3_core.py index 219a65728..7c9aea49e 100644 --- a/sotodlib/core/g3_core.py +++ b/sotodlib/core/g3_core.py @@ -6,7 +6,7 @@ """ -from spt3g import core +from so3g.spt3g import core class DataG3Module(object): diff --git a/sotodlib/mapmaking/ml_mapmaker.py b/sotodlib/mapmaking/ml_mapmaker.py index 567470f04..06333ec67 100644 --- a/sotodlib/mapmaking/ml_mapmaker.py +++ b/sotodlib/mapmaking/ml_mapmaker.py @@ -1,6 +1,7 @@ import numpy as np import h5py import so3g +from typing import Optional from pixell import bunch, enmap, tilemap from pixell import utils as putils @@ -13,7 +14,7 @@ class MLMapmaker: - def __init__(self, signals=[], noise_model=None, dtype=np.float32, verbose=False): + def __init__(self, signals=[], noise_model=None, dtype=np.float32, verbose=False, glitch_flags:str = "flags.glitch_flags"): """Initialize a Maximum Likelihood Mapmaker. Arguments: * signals: List of Signal-objects representing the models that will be solved @@ -33,6 +34,7 @@ def __init__(self, signals=[], noise_model=None, dtype=np.float32, verbose=False self.data = [] self.dof = MultiZipper() self.ready = False + self.glitch_flags_path = glitch_flags def add_obs(self, id, obs, deslope=True, noise_model=None, signal_estimate=None): # Prepare our tod @@ -67,7 +69,7 @@ def add_obs(self, id, obs, deslope=True, noise_model=None, signal_estimate=None) tod = nmat.apply(tod) # Add the observation to each of our signals for signal in self.signals: - signal.add_obs(id, obs, nmat, tod) + signal.add_obs(id, obs, nmat, tod, glitch_flags=self.glitch_flags_path) # Save what we need about this observation self.data.append(bunch.Bunch(id=id, ndet=obs.dets.count, nsamp=len(ctime), dets=obs.dets.vals, nmat=nmat)) @@ -169,7 +171,7 @@ def __init__(self, name, ofmt, output, ext, glitch_flags: str = "flags.glitch_fl self.dof = None self.ready = False self.glitch_flags = glitch_flags - def add_obs(self, id, obs, nmat, Nd): pass + def add_obs(self, id, obs, nmat, Nd, glitch_flags:Optional[str]): pass def prepare(self): self.ready = True def forward (self, id, tod, x): pass def backward(self, id, tod, x): pass @@ -210,7 +212,7 @@ def __init__(self, shape, wcs, comm, comps="TQU", name="sky", ofmt="{name}", out self.div = enmap.zeros((ncomp,ncomp)+shape, wcs, dtype=dtype) self.hits= enmap.zeros( shape, wcs, dtype=dtype) - def add_obs(self, id, obs, nmat, Nd, pmap=None): + def add_obs(self, id, obs, nmat, Nd, pmap=None, glitch_flags: Optional[str] = None): """Add and process an observation, building the pointing matrix and our part of the RHS. "obs" should be an Observation axis manager, nmat a noise model, representing the inverse noise covariance matrix, @@ -218,7 +220,8 @@ def add_obs(self, id, obs, nmat, Nd, pmap=None): """ Nd = Nd.copy() # This copy can be avoided if build_obs is split into two parts ctime = obs.timestamps - pcut = PmatCut(get_flags_from_path(obs, self.glitch_flags)) # could pass this in, but fast to construct + gflags = glitch_flags if glitch_flags is not None else self.glitch_flags + pcut = PmatCut(get_flags_from_path(obs, gflags)) # could pass this in, but fast to construct if pmap is None: # Build the local geometry and pointing matrix for this observation if self.recenter: @@ -381,12 +384,14 @@ def __init__(self, comm, name="cut", ofmt="{name}_{rank:02}", dtype=np.float32, self.rhs = [] self.div = [] - def add_obs(self, id, obs, nmat, Nd): + def add_obs(self, id, obs, nmat, Nd, glitch_flags: Optional[str] = None): """Add and process an observation. "obs" should be an Observation axis manager, nmat a noise model, representing the inverse noise covariance matrix, and Nd the result of applying the noise model to the detector time-ordered data.""" Nd = Nd.copy() # This copy can be avoided if build_obs is split into two parts - pcut = PmatCut(get_flags_from_path(obs, self.glitch_flags), model=self.cut_type) + + gflags = glitch_flags if glitch_flags is not None else self.glitch_flags + pcut = PmatCut(get_flags_from_path(obs, gflags), model=self.cut_type) # Build our RHS obs_rhs = np.zeros(pcut.njunk, self.dtype) pcut.backward(Nd, obs_rhs) @@ -450,7 +455,7 @@ def translate(self, other, junk): so3g.translate_cuts(odata.pcut.cuts, sdata.pcut.cuts, sdata.pcut.model, sdata.pcut.params, junk[odata.i1:odata.i2], res[sdata.i1:sdata.i2]) return res - def transeval(self, id, obs, other, junk, tod): + def transeval(self, id, obs, other, junk, tod, glitch_flags: Optional[str] = None): """Translate data junk from SignalCut other to the current SignalCut, and then evaluate it for the given observation, returning a tod. This is used when building a signal-free tod for the noise model @@ -458,7 +463,8 @@ def transeval(self, id, obs, other, junk, tod): self._checkcompat(other) # We have to make a pointing matrix from scratch because add_obs # won't have been called yet at this point - spcut = PmatCut(get_flags_from_path(obs, self.glitch_flags), model=self.cut_type) + gflags = glitch_flags if glitch_flags is not None else self.glitch_flags + spcut = PmatCut(get_flags_from_path(obs, gflags), model=self.cut_type) # We do have one for other though, since that will be the output # from the previous round of multiplass mapmaking. odata = other.data[id] diff --git a/sotodlib/mapmaking/utilities.py b/sotodlib/mapmaking/utilities.py index b27a78037..423cf7a34 100644 --- a/sotodlib/mapmaking/utilities.py +++ b/sotodlib/mapmaking/utilities.py @@ -1,4 +1,4 @@ -from typing import Any, Union +from typing import Any, Union, Optional import numpy as np import so3g @@ -441,8 +441,48 @@ def rangemat_sum(rangemat): res[i] = np.sum(ra[:,1]-ra[:,0]) return res -def find_usable_detectors(obs, maxcut=0.1): - ncut = rangemat_sum(obs.flags.glitch_flags) +def flags_in_path( + aman: core.AxisManager, rpath: str, sep: str = "." +) -> bool: + """ + This function allows to pull data from an AxisManager based on a path. + Parameters: + - aman: An Axis Manager object + - path: a string with a recursive path to extract data. The path is separated via a sep. + For example 'flags.glitch_flags' + - sep: separator. Defaults to `.` + """ + + rpath = rpath.split(sep=sep) + flags = aman.copy() + while rpath and flags is not None: + path = rpath.pop() + flags = flags[path] + + return flags is not None + + +def get_flags_from_path( + aman: core.AxisManager, rpath: str, sep: str = "." +) -> Union[so3g.proj.RangesMatrix, Any]: + """ + This function allows to pull data from an AxisManager based on a path. + Parameters: + - aman: An Axis Manager object + - path: a string with a recursive path to extract data. The path is separated via a sep. + For example 'flags.glitch_flags' + - sep: separator. Defaults to `.` + """ + + flags = aman.copy() + for path in rpath.split(sep=sep): + flags = flags[path] + + return flags + + +def find_usable_detectors(obs, maxcut=0.1, glitch_flags: str = "flags.glitch_flags"): + ncut = rangemat_sum(get_flags_from_path(obs, glitch_flags)) good = ncut < obs.samps.count * maxcut return obs.dets.vals[good] @@ -501,7 +541,7 @@ def downsample_obs(obs, down): if isinstance(val, core.AxisManager): res.wrap(key, val) else: - axdesc = [(k,v) for k,v in enumerate(axes) if v is not None] + axdesc = [(k, v) for k, v in enumerate(axes) if v is not None] res.wrap(key, val, axdesc) # The normal sample stuff res.wrap("timestamps", obs.timestamps[::down], [(0, "samps")]) @@ -513,33 +553,23 @@ def downsample_obs(obs, down): # The cuts # obs.flags will contain all types of flags. We should query it for glitch_flags and source_flags - cut_keys = ["glitch_flags"] + cut_keys = [] + if flags_in_path(obs, "glitch_flags"): + cut_keys.append("glitch_flags") + elif flags_in_path(obs, "flags.glitch_flags"): + cut_keys.append("flags.glitch_flags") - if "source_flags" in obs.flags: + if flags_in_path(obs, "source_flags"): cut_keys.append("source_flags") + elif flags_in_path(obs, "flags.source_flags"): + cut_keys.append("flags.source_flags") # We need to add a res.flags FlagManager to res res = res.wrap('flags', core.FlagManager.for_tod(res)) for key in cut_keys: - res.flags.wrap(key, downsample_cut(getattr(obs.flags, key), down), [(0,"dets"),(1,"samps")]) + res.flags.wrap(key, downsample_cut(get_flags_from_path(obs, key), down), [(0,"dets"),(1,"samps")]) # Not sure how to deal with flags. Some sort of or-binning operation? But it # doesn't matter anyway return res - -def get_flags_from_path(aman: core.AxisManager, rpath: str, sep: str=".") -> Union[so3g.proj.RangesMatrix, Any]: - """ - This function allows to pull data from an AxisManager based on a path. - Parameters: - - aman: An Axis Manager object - - path: a string with a recursive path to extract data. The path is separated via a sep. - For example 'flags.glitch_flags' - - sep: separator. Defaults to `.` - """ - - tmp = aman.copy() - for path in rpath.split(sep=sep): - tmp = tmp[path] - - return tmp