diff --git a/src/toast/intervals.py b/src/toast/intervals.py index 3262db714..a0706e4ba 100644 --- a/src/toast/intervals.py +++ b/src/toast/intervals.py @@ -262,17 +262,38 @@ def __and__(self, other): result = list() curself = 0 curother = 0 + lenself = len(self.data) + lenother = len(other) + + # Find the first entry in ivals whose last is not smaller than val, + # searching within the range of indices [ileft,iright) + def lbound(ivals, val, ileft, iright): + # represents the width of the remaining interval + count = iright-ileft + while count>0: + i = ileft + count//2 + if ivals[i].last < val: + ileft = i+1 + count -= 1 + count//2 + else: + count = count//2 + return ileft # Walk both sequences, building up the intersection. - while (curself < len(self.data)) and (curother < len(other)): + while (curself < lenself) and (curother < lenother): low = max(self.data[curself].first, other[curother].first) high = min(self.data[curself].last, other[curother].last) if low <= high: result.append((self.timestamps[low], self.timestamps[high], low, high)) - if self.data[curself].last < other[curother].last: - curself += 1 + if self.data[curself].last < other[curother].last: + curself += 1 + else: + curother += 1 else: - curother += 1 + if self.data[curself].last < other[curother].last: + curself = lbound(self.data, other[curother].first, curself, lenself) + else: + curother = lbound(other, self.data[curself].first, curother, lenother) return IntervalList( self.timestamps, diff --git a/src/toast/ops/sim_ground.py b/src/toast/ops/sim_ground.py index fbdb0676d..22ee224da 100644 --- a/src/toast/ops/sim_ground.py +++ b/src/toast/ops/sim_ground.py @@ -796,7 +796,7 @@ def _exec(self, data, detectors=None, **kwargs): if self.det_data is not None: exists_data = ob.detdata.ensure( - self.det_data, dtype=np.float64, detectors=dets + self.det_data, dtype=np.float64, detectors=dets, units=u.Kelvin ) if self.det_flags is not None: diff --git a/src/toast/ops/sim_tod_noise.py b/src/toast/ops/sim_tod_noise.py index 2991146bb..f7e045345 100644 --- a/src/toast/ops/sim_tod_noise.py +++ b/src/toast/ops/sim_tod_noise.py @@ -277,7 +277,7 @@ def _exec(self, data, detectors=None, **kwargs): # detectors within the observation... # Make sure correct output exists - exists = ob.detdata.ensure(self.det_data, detectors=dets) + exists = ob.detdata.ensure(self.det_data, detectors=dets, units=u.Kelvin) # Get the sample rate from the data. We also have nominal sample rates # from the noise model and also from the focalplane. Perhaps we should add diff --git a/src/toast/spt3g/spt3g_export.py b/src/toast/spt3g/spt3g_export.py index d34691329..880a43107 100644 --- a/src/toast/spt3g/spt3g_export.py +++ b/src/toast/spt3g/spt3g_export.py @@ -13,6 +13,7 @@ from ..intervals import IntervalList from ..timing import function_timer from ..utils import Environment, Logger, object_fullname +from .. import qarray as qa from .spt3g_utils import ( available, compress_timestream, @@ -26,6 +27,7 @@ if available: from spt3g import core as c3g + from spt3g import calibration as spt_cal @function_timer @@ -63,6 +65,29 @@ def export_shared(obs, name, view_name=None, view_index=0, g3t=None): return g3t(sview.flatten().tolist()) +def export_focalplane(fplane): + out_props = spt_cal.BolometerPropertiesMap() + zaxis = np.array([0.0, 0.0, 1.0], dtype=np.float64) + for idx,det_name in enumerate(fplane.keys()): + bolo = spt_cal.BolometerProperties() + + dir = fplane.detector_data["quat"][idx] + rdir = qa.rotate(dir, zaxis).flatten() + ang = np.arctan2(rdir[1], rdir[0]) + + mag = np.arccos(rdir[2]) * c3g.G3Units.rad + + bolo.physical_name = det_name # is this a 'physical' name? + bolo.x_offset = mag * np.cos(ang) + bolo.y_offset = mag * np.sin(ang) + bolo.band = (fplane.detector_data["bandcenter"][idx] / u.GHz).value * c3g.G3Units.GHz + bolo.pol_angle = (fplane.detector_data["pol_angle"][idx] / u.rad).value * c3g.G3Units.rad + bolo.pol_efficiency = fplane.detector_data["pol_efficiency"][idx] + + out_props[det_name] = bolo + return out_props + + @function_timer def export_detdata( obs, name, view_name=None, view_index=0, g3t=None, times=None, compress=None @@ -290,6 +315,9 @@ def __call__(self, obs): # Construct calibration frame cal = c3g.G3Frame(c3g.G3FrameType.Calibration) + # Output focal plane using proper SPT structure + cal["BolometerProperties"] = export_focalplane(obs.telescope.focalplane) + # Serialize focalplane to HDF5 bytes and write to frame. byte_writer = io.BytesIO() with h5py.File(byte_writer, "w") as f: @@ -302,8 +330,12 @@ def __call__(self, obs): byte_writer = io.BytesIO() with h5py.File(byte_writer, "w") as f: obs[m_in].save_hdf5(f, comm=None, force_serial=True) + if m_out in cal: + del cal[m_out] cal[m_out] = c3g.G3VectorUnsignedChar(byte_writer.getvalue()) del byte_writer + if f"{m_out}_class" in cal: + del cal[f"{m_out}_class"] cal[f"{m_out}_class"] = c3g.G3String(object_fullname(obs[m_in].__class__)) return ob, cal @@ -357,6 +389,7 @@ def __init__( frame_intervals=None, shared_names=list(), det_names=list(), + split_interval_names=list(), interval_names=list(), compress=False, ): @@ -364,6 +397,7 @@ def __init__( self._frame_intervals = frame_intervals self._shared_names = shared_names self._det_names = det_names + self._split_interval_names = split_interval_names self._interval_names = interval_names self._compress = compress @@ -376,26 +410,63 @@ def __call__(self, obs): log = Logger.get() frame_intervals = self._frame_intervals if frame_intervals is None: - # We are using the sample set distribution for our frame boundaries. frame_intervals = "frames" timespans = list() - offset = 0 - n_frames = 0 - first_set = obs.dist.samp_sets[obs.comm.group_rank].offset - n_set = obs.dist.samp_sets[obs.comm.group_rank].n_elem - for sset in range(first_set, first_set + n_set): - for chunk in obs.dist.sample_sets[sset]: - timespans.append( - ( - obs.shared[self._timestamp_names[0]][offset], - obs.shared[self._timestamp_names[0]][offset + chunk - 1], + + if len(self._split_interval_names) > 0: + # Split up data into frames by types of interval, as is normal for observations. + # There may be several types of intervals listed in _interval_names, + # and we need to pull them out in order + # So, treat each obs.intervals[ivl_key] as a queue of intervals, + # and cycle through them pulling out the first from each queue + # to add to the list of frame intervals + def earlier(i1, i2): + if i2 is None: + return True + return i1.start < i2.start + queue_indices = [0] * len(self._split_interval_names) + done = False + while not done: + done = True + next = None + nextIdx = -1 + idx = 0 + for ivl_key in self._split_interval_names: + if queue_indices[idx] < len(obs.intervals[ivl_key]): + done = False # if any queue has at least one more interval, we must continue + if earlier(obs.intervals[ivl_key][queue_indices[idx]], next): + next = obs.intervals[ivl_key][queue_indices[idx]] + nextIdx = idx + idx += 1 + if not done: + timespans.append( + ( + obs.intervals[self._split_interval_names[nextIdx]][queue_indices[nextIdx]].start, + obs.intervals[self._split_interval_names[nextIdx]][queue_indices[nextIdx]].stop, + ) ) - ) - n_frames += 1 - offset += chunk - obs.intervals.create_col( - frame_intervals, timespans, obs.shared[self._timestamp_names[0]] - ) + queue_indices[nextIdx] += 1 + obs.intervals.create_col(frame_intervals, timespans, obs.shared[self._timestamp_names[0]]) + else: + # Divide data into frames arbitrarily, which may be unsuitable for further processing + # We are using the sample set distribution for our frame boundaries. + offset = 0 + n_frames = 0 + first_set = obs.dist.samp_sets[obs.comm.group_rank].offset + n_set = obs.dist.samp_sets[obs.comm.group_rank].n_elem + for sset in range(first_set, first_set + n_set): + for chunk in obs.dist.sample_sets[sset]: + timespans.append( + ( + obs.shared[self._timestamp_names[0]][offset], + obs.shared[self._timestamp_names[0]][offset + chunk - 1], + ) + ) + n_frames += 1 + offset += chunk + obs.intervals.create_col( + frame_intervals, timespans, obs.shared[self._timestamp_names[0]] + ) output = list() frame_view = obs.view[frame_intervals] @@ -456,7 +527,7 @@ def __call__(self, obs): if len(self._interval_names) > 0: tview = obs.view[frame_intervals].shared[self._timestamp_names[0]][ivw] iframe = IntervalList( - obs.shared[self._timestamp_names[0]], + np.array(obs.shared[self._timestamp_names[0]]), timespans=[(tview[0], tview[-1])], ) for ivl_key, ivl_val in self._interval_names: @@ -465,6 +536,11 @@ def __call__(self, obs): ivl_key, iframe, ) + # Tag frames which contain a period of turning the telescope around + # This may not work usefully if the frame splitting was not based on + # suitable intervals + if ivl_val == "intervals_turnaround" and len(frame[ivl_val])>0: + frame["Turnaround"] = c3g.G3Bool(True) output.append(frame) # Delete our temporary frame interval if we created it if self._frame_intervals is None: diff --git a/src/toast/spt3g/spt3g_utils.py b/src/toast/spt3g/spt3g_utils.py index bcabb4e81..274d9f5c9 100644 --- a/src/toast/spt3g/spt3g_utils.py +++ b/src/toast/spt3g/spt3g_utils.py @@ -6,6 +6,7 @@ import numpy as np from astropy import units as u +from scipy.spatial.transform import Rotation from ..utils import Environment, Logger @@ -95,8 +96,9 @@ def to_g3_map_array_type(dt): def to_g3_unit(aunit): """Convert astropy unit to G3 timestream unit. - We convert our input units to SI base units (no milli-, micro-, etc prefix). - We also return the scale factor needed to transform to this unit. + We convert our input units to SPT-3G base units. + We also return the scale factor needed to transform to this unit, + defined as the ratio between definitions of the fundmental SI units. Args: aunit (astropy.unit): The input unit. @@ -111,47 +113,47 @@ def to_g3_unit(aunit): scale = 1.0 * aunit # Try to convert the units to the various types of quantities try: - scale = scale.to(u.volt) + scale = scale.to(u.volt)*c3g.G3Units.volt return (c3g.G3TimestreamUnits.Voltage, scale.value) except Exception: pass try: - scale = scale.to(u.watt) + scale = scale.to(u.watt)*c3g.G3Units.watt return (c3g.G3TimestreamUnits.Power, scale.value) except Exception: pass try: - scale = scale.to(u.ohm) + scale = scale.to(u.ohm)*(c3g.G3Units.volt/c3g.G3Units.amp) return (c3g.G3TimestreamUnits.Resistance, scale.value) except Exception: pass try: - scale = scale.to(u.ampere) + scale = scale.to(u.ampere)*c3g.G3Units.ampere return (c3g.G3TimestreamUnits.Current, scale.value) except Exception: pass try: - scale = scale.to(u.meter) + scale = scale.to(u.meter)*c3g.G3Units.meter return (c3g.G3TimestreamUnits.Distance, scale.value) except Exception: pass try: - scale = scale.to(u.pascal) + scale = scale.to(u.bar)*c3g.G3Units.bar return (c3g.G3TimestreamUnits.Pressure, scale.value) except Exception: pass try: - scale = scale.to(u.radian) + scale = scale.to(u.radian)*c3g.G3Units.radian return (c3g.G3TimestreamUnits.Angle, scale.value) except Exception: pass try: - scale = scale.to(u.Jy) + scale = scale.to(u.Jy)*c3g.G3Units.Jy return (c3g.G3TimestreamUnits.FluxDensity, scale.value) except Exception: pass try: - scale = scale.to(u.Kelvin) + scale = scale.to(u.Kelvin)*c3g.G3Units.kelvin return (c3g.G3TimestreamUnits.Tcmb, scale.value) except Exception: pass @@ -161,7 +163,7 @@ def to_g3_unit(aunit): def from_g3_unit(gunit): """Convert G3 timestream unit to astropy. - This function assumes that the quantities are in SI units. + This function assumes that the quantities are in SPT-3G units. Args: gunit (G3TimestreamUnit): The input units. @@ -173,23 +175,23 @@ def from_g3_unit(gunit): if gunit == c3g.G3TimestreamUnits.Counts: return u.dimensionless_unscaled elif gunit == c3g.G3TimestreamUnits.Voltage: - return u.volt + return u.def_unit('SPT_Voltage', u.volt/c3g.G3Units.volt) elif gunit == c3g.G3TimestreamUnits.Power: - return u.watt + return u.def_unit('SPT_Power', u.watt/c3g.G3Units.watt) elif gunit == c3g.G3TimestreamUnits.Resistance: - return u.ohm + return u.def_unit('SPT_Resistance', u.ohm/(c3g.G3Units.volt/c3g.G3Units.amp)) elif gunit == c3g.G3TimestreamUnits.Current: - return u.ampere + return u.def_unit('SPT_Current', u.ampere/c3g.G3Units.ampere) elif gunit == c3g.G3TimestreamUnits.Distance: - return u.meter + return u.def_unit('SPT_Distance', u.meter/c3g.G3Units.meter) elif gunit == c3g.G3TimestreamUnits.Pressure: - return u.pascal + return u.def_unit('SPT_Pressure', u.bar/c3g.G3Units.bar) elif gunit == c3g.G3TimestreamUnits.Angle: - return u.radian + return u.def_unit('SPT_Angle', u.radian/c3g.G3Units.radian) elif gunit == c3g.G3TimestreamUnits.FluxDensity: - return u.Jy + return u.def_unit('SPT_FluxDensity', u.Jy/c3g.G3Units.Jy) elif gunit == c3g.G3TimestreamUnits.Tcmb: - return u.Kelvin + return u.def_unit('SPT_Tcmb', u.Kelvin/c3g.G3Units.kelvin) else: return u.dimensionless_unscaled @@ -329,6 +331,28 @@ def to_g3_time(input): return c3g.G3VectorTime([c3g.G3Time(x * 1e8) for x in input]) +from_3g_rotation = Rotation.from_rotvec([0, np.pi/2., 0]) + + +def from_g3_quat(q: c3g.quat): + """Convert a single SPT-3G quaternion to TOAST conventions. + + This reorders the components, and rotates from the boresight being aligned + with the x-axis to the z-axis being so aligned. + TODO: Is such a rotation appropriate for all quaternions being imported? + + Args: + q (spt3g.core.quat): An SPT-3G quaternion + + Returns: + (array): A TOAST quaternion + """ + # roll components from standard order to scipy order, + # then perform the rotation to change coordinate systems + r = Rotation.from_quat([q.b, q.c, q.d, q.a]) + return (r * from_3g_rotation).as_quat() + + def from_g3_quats(input): """Convert Spt3G quaternions into TOAST format. @@ -342,10 +366,33 @@ def from_g3_quats(input): """ if isinstance(input, c3g.G3VectorQuat): - return np.array([np.array([x.b, x.c, x.d, x.a]) for x in input]) + return np.array([from_g3_quat(x) for x in input]) else: # Assume it is a scalar - return np.array([input.b, input.c, input.d, input.a]) + return from_g3_quat(input) + + +to_3g_rotation = Rotation.from_rotvec([0, -np.pi/2., 0]) + + +def to_g3_quat(q): + """Convert a single TOAST quaternion to SPT-3G conventions. + + This corrects storage order, and rotates from the z-axis being aligned with + the boresight to the x-axis being so aligned. + TODO: Is such a rotation appropriate for all quaternions being exported? + + Args: + q (array): A TOAST quaternion + + Returns: + (spt3g.core.quat): An SPT-3G quaternion + """ + # first, perform rotation while in TOAST format, then convert back to an + # explicit quaternion to roll components to standard order + r = Rotation.from_quat(q) + q = (r * to_3g_rotation).as_quat() + return c3g.quat(q[3], q[0], q[1], q[2]) def to_g3_quats(input): @@ -362,10 +409,10 @@ def to_g3_quats(input): """ if len(input.shape) == 2: # Array of values - return c3g.G3VectorQuat([c3g.quat(x[3], x[0], x[1], x[2]) for x in input]) + return c3g.G3VectorQuat([to_g3_quat(q) for q in input]) else: # One value - return c3g.quat(input[3], input[0], input[1], input[2]) + return to_g3_quat(input) # Callable that just builds a list of frames diff --git a/workflows/toast_sim_ground.py b/workflows/toast_sim_ground.py index 619d3de01..0c1852508 100644 --- a/workflows/toast_sim_ground.py +++ b/workflows/toast_sim_ground.py @@ -628,23 +628,24 @@ def dump_spt3g(job, args, data): det_names=[ ( ops.sim_noise.det_data, - ops.sim_noise.det_data, + "CalTimestreams", c3g.G3TimestreamMap, ), # ("flags", "detector_flags", c3g.G3TimestreamMap), ], + # split data into frames where it is covered by these intervals + split_interval_names=[ + ops.sim_ground.scan_leftright_interval, + ops.sim_ground.turn_leftright_interval, + ops.sim_ground.scan_rightleft_interval, + ops.sim_ground.turn_rightleft_interval, + ops.sim_ground.elnod_interval + ], + # export information about these intervals into frames interval_names=[ - (ops.sim_ground.scan_leftright_interval, "intervals_scan_leftright"), - (ops.sim_ground.turn_leftright_interval, "intervals_turn_leftright"), - (ops.sim_ground.scan_rightleft_interval, "intervals_scan_rightleft"), - (ops.sim_ground.turn_rightleft_interval, "intervals_turn_rightleft"), - (ops.sim_ground.elnod_interval, "intervals_elnod"), - (ops.sim_ground.scanning_interval, "intervals_scanning"), (ops.sim_ground.turnaround_interval, "intervals_turnaround"), - (ops.sim_ground.sun_up_interval, "intervals_sun_up"), - (ops.sim_ground.sun_close_interval, "intervals_sun_close"), ], - compress=True, + compress=False, # Compression must be off for data to be intelligible to spt3g_software ) exporter = t3g.export_obs( meta_export=meta_exporter,