diff --git a/sotodlib/mapmaking/demod_mapmaker.py b/sotodlib/mapmaking/demod_mapmaker.py index 7511044d9..54e5b0b35 100644 --- a/sotodlib/mapmaking/demod_mapmaker.py +++ b/sotodlib/mapmaking/demod_mapmaker.py @@ -5,11 +5,12 @@ you accumulate observations into the div and rhs maps. For examples how to use look at docstring of DemodMapmaker. """ -__all__ = ['DemodMapmaker','DemodSignal','DemodSignalMap'] -import numpy as np -from pixell import enmap, utils, tilemap, bunch +__all__ = ['DemodMapmaker','DemodSignal','DemodSignalMap','make_demod_map'] +import numpy as np, os +from pixell import enmap, utils as putils, tilemap, bunch, mpi import so3g.proj +from .. import core from .. import coords from .utilities import recentering_to_quat_lonlat, evaluate_recentering, MultiZipper, unarr, safe_invert_div from .utilities import import_optional, get_flags @@ -213,7 +214,7 @@ def for_rectpix(cls, shape, wcs, comm, comps="TQU", name="sky", ofmt="{name}", o ext : str, optional The extension used for the files. dtype : numpy.dtype - The data type to use for the time-ordered data. + The data type to use for the maps. sys : str or None, optional The coordinate system to map. Defaults to equatorial recenter : str or None @@ -338,22 +339,14 @@ def prepare(self): self.hits = tilemap.redistribute(self.hits,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.ready = True @property def ncomp(self): return len(self.comps) - def to_work(self, map): - if self.tiled: return tilemap.redistribute(map, self.comm, self.geo_work.active) - else: return map.copy() - - def from_work(self, map): - if self.tiled: return tilemap.redistribute(map, self.comm, self.rhs.geometry.active) - else: return utils.allreduce(map, self.comm) - def write(self, prefix, tag, m): if not self.output: return oname = self.ofmt.format(name=self.name) @@ -385,6 +378,207 @@ def write(self, prefix, tag, m): return oname +def setup_demod_map(noise_model, shape=None, wcs=None, nside=None, + comm=mpi.COMM_WORLD, comps='TQU', split_labels=None, + singlestream=False, dtype_tod=np.float32, + dtype_map=np.float32, recenter=None, verbose=0): + """ + Setup the classes for demod mapmaking and return + a DemodMapmmaker object + """ + if shape is not None and wcs is not None: + if split_labels==None: + # this is the case where we did not request any splits at all + signal_map = DemodSignalMap.for_rectpix(shape, wcs, comm, comps=comps, + dtype=dtype_map, tiled=False, + ofmt="", singlestream=singlestream, + recenter=recenter ) + else: + # this is the case where we asked for at least 2 splits (1 split set). + # We count how many split we'll make, we need to define the Nsplits + # maps inside the DemodSignalMap + Nsplits = len(split_labels) + signal_map = DemodSignalMap.for_rectpix(shape, wcs, comm, comps=comps, + dtype=dtype_map, tiled=False, + ofmt="", Nsplits=Nsplits, + singlestream=singlestream, + recenter=recenter) + elif nside is not None: + if split_labels==None: + # this is the case where we did not request any splits at all + signal_map = DemodSignalMap.for_healpix(nside, nside_tile='auto', + comps=comps, dtype=dtype_map, + ofmt="", singlestream=singlestream, + ext="fits.gz") + else: + # this is the case where we asked for at least 2 splits (1 split set). + # We count how many split we'll make, we need to define the Nsplits + # maps inside the DemodSignalMap + Nsplits = len(split_labels) + signal_map = DemodSignalMap.for_healpix(nside, nside_tile='auto', + comps=comps, dtype=dtype_map, + ofmt="", Nsplits=Nsplits, + singlestream=singlestream, + ext="fits.gz") + signals = [signal_map] + mapmaker = DemodMapmaker(signals, noise_model=noise_model, + dtype=dtype_tod, + verbose=verbose>0, + singlestream=singlestream) + return mapmaker + +def write_demod_maps(prefix, data, split_labels=None): + """ + Write maps from data into files + """ + if split_labels==None: + # we have no splits, so we save index 0 of the lists + data.signal.write(prefix, "full_wmap", data.wmap[0]) + data.signal.write(prefix, "full_weights", data.weights[0]) + data.signal.write(prefix, "full_hits", data.signal.hits) + else: + # we have splits + Nsplits = len(split_labels) + for n_split in range(Nsplits): + data.signal.write(prefix, "%s_wmap"%split_labels[n_split], + data.wmap[n_split]) + data.signal.write(prefix, "%s_weights"%split_labels[n_split], + data.weights[n_split]) + data.signal.write(prefix, "%s_hits"%split_labels[n_split], + data.signal.hits[n_split]) + +def write_demod_info(oname, info, split_labels=None): + putils.mkdir(os.path.dirname(oname)) + if split_labels==None: + bunch.write(oname+'_full_info.hdf', info[0]) + else: + # we have splits + Nsplits = len(split_labels) + for n_split in range(Nsplits): + bunch.write(oname+'_%s_info.hdf'%split_labels[n_split], info[n_split]) + +def make_demod_map(context, obslist, noise_model, info, + preprocess_config, prefix, shape=None, wcs=None, + nside=None, comm=mpi.COMM_WORLD, comps="TQU", t0=0, + dtype_tod=np.float32, dtype_map=np.float32, + tag="", verbose=0, split_labels=None, L=None, + site='so_sat3', recenter=None, singlestream=False): + """ + Make a demodulated map from the list of observations in obslist. + + Arguments + --------- + context : str + File path to context used to load obs from. + obslist : dict + The obslist which is the output of the + mapmaking.obs_grouping.build_obslists, contains the information of the + single or multiple obs to map. + noise_model : sotodlib.mapmaking.Nmat + Noise model to pass to DemodMapmaker. + info : list + Information for the database, will be written as a .hdf file. + preprocess_config : dict + Dictionary with the config yaml file for the preprocess database. + prefix : str + Prefix for the output files + shape : tuple, optional + Shape of the geometry to use for mapping. + wcs : dict, optional + WCS kernel of the geometry to use for mapping. + nside : int, optional + Nside for healpix pixelization + comps : str, optional + Which components to map, only TQU supported for now. + t0 : int, optional + Ctime to use as the label in the map files. + dtype_tod : numpy.dtype, optional + The data type to use for the time-ordered data. Only tested + with float32. + dtype_map : numpy.dtype, optional + The data type to use for the maps. + tag : str, optional + Prefix tag for the logger. + verbose : bool, optional + split_labels : list or None, optional + A list of strings with the splits requested. If None then no splits + were asked for, i.e. we will produce one map. + L : logger, optional + Logger for printing on the screen. + site : str, optional + Platform name for the pointing matrix. + recenter : str or None + String to make object-centered maps, such as Moon/Sun/Planet centered maps. + Look at sotodlib.mapmaking.parse_recentering for details. + singlestream : Bool + If True, do not perform demodulated filter+bin mapmaking but + rather regular filter+bin mapmaking, i.e. map from obs.signal + rather than from obs.dsT, obs.demodQ, obs.demodU. + + Returns + ------- + errors : list + List of errors from preprocess database. To be used in cleanup_mandb. + outputs : list + List of outputs from preprocess database. To be used in cleanup_mandb. + """ + from .. import site_pipeline + context = core.Context(context) + if L is None: + L = site_pipeline.util.init_logger("Demod filterbin mapmaking") + pre = "" if tag is None else tag + " " + if comm.rank == 0: L.info(pre + "Initializing equation system") + mapmaker = setup_demod_map(noise_model, shape=shape, wcs=wcs, nside=nside, + comm=comm, comps=comps, split_labels=split_labels, + singlestream=singlestream, dtype_tod=dtype_tod, + dtype_map=dtype_map, recenter=recenter, verbose=verbose) + + if comm.rank == 0: L.info(pre + "Building RHS") + # And feed it with our observations + nobs_kept = 0 + errors = [] ; outputs = []; # PENDING: do an allreduce of these. + # not needed for atomic maps, but needed for + # depth-1 maps + for oi in range(len(obslist)): + obs_id, detset, band = obslist[oi][:3] + name = "%s:%s:%s" % (obs_id, detset, band) + error, output, obs = site_pipeline.preprocess_tod.preproc_or_load_group(obs_id, + configs=preprocess_config, + dets={'wafer_slot':detset, 'wafer.bandpass':band}, + logger=L, context=context, overwrite=False) + errors.append(error) ; outputs.append(output) ; + if error not in [None,'load_success']: + L.info('tod %s:%s:%s failed in the prepoc database'%(obs_id,detset,band)) + continue + obs.wrap("weather", np.full(1, "toco")) + obs.wrap("site", np.full(1, site)) + obs.flags.wrap('glitch_flags', obs.preprocess.turnaround_flags.turnarounds + + obs.preprocess.jumps_2pi.jump_flag + obs.preprocess.glitches.glitch_flags, ) + mapmaker.add_obs(name, obs, split_labels=split_labels) + L.info('Done with tod %s:%s:%s'%(obs_id,detset,band)) + nobs_kept += 1 + nobs_kept = comm.allreduce(nobs_kept) + # if we skip all the obs then we return error and output + if nobs_kept == 0: + return errors, outputs + + for signal in mapmaker.signals: + signal.prepare() + if comm.rank == 0: L.info(pre + "Writing F+B outputs") + wmap = [] + weights = [] + # mapmaker.signals[0] is signal_map + for n_split in range(mapmaker.signals[0].Nsplits): + wmap.append( mapmaker.signals[0].rhs[n_split] ) + div = np.diagonal(mapmaker.signals[0].div[n_split], axis1=0, axis2=1) + div = np.moveaxis(div, -1, 0) # this moves the last axis to the 0th position + weights.append(div) + mapdata = bunch.Bunch(wmap=wmap, weights=weights, signal=mapmaker.signals[0], t0=t0) + # output to files + write_demod_maps(prefix, mapdata, split_labels=split_labels, ) + write_demod_info(prefix, info, split_labels=split_labels ) + return errors, outputs + def project_rhs_demod(pmap, signalT, signalQ, signalU, det_weightsT, det_weightsQU, wrapper=lambda x:x): """ Project demodulated T, Q, U timestreams into weighted maps. diff --git a/sotodlib/mapmaking/obs_grouping.py b/sotodlib/mapmaking/obs_grouping.py index 92fc68ff5..ac8ee329c 100644 --- a/sotodlib/mapmaking/obs_grouping.py +++ b/sotodlib/mapmaking/obs_grouping.py @@ -15,13 +15,17 @@ depth-1 maps, etc. """ -__all__ = ['build_obslists'] +__all__ = ['build_obslists','NoTODFound'] import numpy as np from pixell import utils from scipy import ndimage from .utilities import get_ids +class NoTODFound(Exception): + def __init__(self, msg): + self.msg = msg + def build_obslists(context, query, mode=None, nset=None, wafer=None, freq=None, ntod=None, tods=None, fixed_time=None, mindur=None, ): """ @@ -78,8 +82,7 @@ def build_obslists(context, query, mode=None, nset=None, wafer=None, if tods: ids = np.array(eval("ids"+tods)).reshape(-1) if ntod: ids = ids[:ntod] if len(ids) == 0: - print("No tods found!") - sys.exit(1) + raise NoTODFound("No tods found!") # Extract info about the selected ids obs_infos = context.obsdb.query("obs_id in (%s)" % ",".join(["'%s'" % id for id in ids])) obs_infos = obs_infos.asarray().view(np.recarray) @@ -100,8 +103,7 @@ def build_obslists(context, query, mode=None, nset=None, wafer=None, periods = split_periods(periods, 24*3600) # If fixed_time was not set, #then we do 24 hrs by default and it will be the same as depth_1 else: - print("Invalid mode!") - sys.exit(1) + raise NoTODFound("Invalid mode!") # We will make one map per period-detset-band obslists = build_period_obslists(obs_infos, periods, context, nset=nset,