def AddBenchData(f): ''' Add the optical bench positions to the frame. ''' if f.type != core.G3FrameType.GcpSlow: return bench_axes = ['y1', 'y2', 'y3', 'x4', 'x5', 'z6'] benchcom = core.G3TimestreamMap() benchpos = core.G3TimestreamMap() benchzero = core.G3TimestreamMap() for i, key in enumerate(bench_axes): # As of 2017-08-03, SCU time is not trustworthy # start = f['antenna0']['scu']['benchSampleTime'][0][0] # stop = f['antenna0']['scu']['benchSampleTime'][0][-1] # For now, do this bit of evil start = f['antenna0']['tracker']['utc'][0][0] stop = f['antenna0']['tracker']['utc'][0][-1] benchcom[key] = core.G3Timestream(f['antenna0']['scu']['benchExpected'][i]) benchcom[key].start = start benchcom[key].stop = stop benchpos[key] = core.G3Timestream(f['antenna0']['scu']['benchActual'][i]) benchpos[key].start = start benchpos[key].stop = stop benchzero[key] = core.G3Timestream(f['antenna0']['scu']['benchZeros'][i]) benchzero[key].start = start benchzero[key].stop = stop f['BenchPosition'] = benchpos f['BenchCommandedPosition'] = benchcom f['BenchZeros'] = benchzero
def __call__(self, frame): if self.bpmkey in frame: self.bpm = frame[self.bpmkey] if self.input in frame: out = {} if self.bands is not None: for band in self.bands: out[band] = core.G3TimestreamMap() for b, ts in frame[self.input].iteritems(): try: band = self.bpm[b].band except KeyError: continue if band not in out: if self.bands is None and \ not (math.isnan(band) or math.isinf(band)): out[band] = core.G3TimestreamMap() else: continue out[band][b] = ts for band in out.keys(): frame['%s%dGHz' % (self.output_root, int(band/core.G3Units.GHz))] = out[band]
def noise_scan_frames(n_frames=3, n_dets=20, input='signal', n_samps=200, samp_rate=0.005 * core.G3Units.second, t_start=core.G3Time('2020-1-1T00:00:00')): """ Generate a list of frames filled with noise data and nothing else. Args: n_frames (int): number of frames to make n_dets (int): number of detectors per frame input (str): name of G3TimestreamMap for detectors, should be some form of 'signal' n_samps (int): number of samples per detector timestream samp_rate (G3Unit.second): detector sampling rate t_start (G3Time): start time of the set of frames """ frame_list = [] for n in range(n_frames): f = core.G3Frame() f.type = core.G3FrameType.Scan tsm = core.G3TimestreamMap() z = np.zeros((n_samps, )) for d in enumerate_det_id(n_dets): tsm[d] = core.G3Timestream(z) tsm.start = t_start tsm.stop = t_start + n_samps * samp_rate tsm = MakeNoiseData().apply(tsm) f[input] = tsm t_start += n_samps * samp_rate frame_list.append(f) return frame_list
def __call__(self, f): if f.type == FT.Calibration and f['cal_type'] == 'focal_plane': self.focal_plane = f if f.type != FT.Scan: return [f] # As long as we have a focal_plane, we can create signal vectors. if self.focal_plane is None: return [f] f['signal'] = core.G3TimestreamMap() # Determine time samples we will be covering. if self.start_time is None: first = f['vertex_enc_raw'].t[0] * core.G3Units.sec self.start_time = core.G3Time( np.ceil(first / self.tick_step) * self.tick_step) # And we will end before... last = core.G3Time(f['vertex_enc_raw'].t[-1] * core.G3Units.sec) n = int((last.time - self.start_time.time) / self.tick_step) end_time = core.G3Time(self.start_time.time + n * self.tick_step) z = np.zeros(n) for k in self.focal_plane['signal_names']: f['signal'][k] = core.G3Timestream(z) # You can't broadcast-set the start and end time unless the # elements are already populated. f['signal'].start = self.start_time f['signal'].stop = end_time self.start_time = end_time return [f]
def __call__(self, frame): if 'DfMuxTransferFunction' in frame: self.default_tf = frame['DfMuxTransferFunction'] if frame.type == core.G3FrameType.Wiring: self.wiringmap = frame['WiringMap'] self.system = frame['ReadoutSystem'] self.convfactors = {} return if frame.type == core.G3FrameType.Housekeeping: self.hkmap = frame['DfMuxHousekeeping'] self.convfactors = {} return if self.keepconversions and frame.type == core.G3FrameType.Observation: self.convfactors = {} self.hkmap = None return if frame.type != core.G3FrameType.Scan: return # Housekeeping data can also be in Scan frames if 'DfMuxHousekeeping' in frame: if self.hkmap is None or (frame['DfMuxHousekeeping'].values()[0].timestamp != self.hkmap.values()[0].timestamp and not self.keepconversions): # XXX: finer-grained check for same values? self.hkmap = frame['DfMuxHousekeeping'] self.convfactors = {} # Now the meat newts = core.G3TimestreamMap() oldts = frame[self.input] for bolo,ts in oldts.items(): if ts.units not in self.convfactors: self.convfactors[ts.units] = {} if ts.units == self.units: convfactor = 1. elif bolo in self.convfactors[ts.units]: convfactor = self.convfactors[ts.units][bolo] else: tf = self.default_tf # XXX: might get from HK data try: convfactor = get_timestream_unit_conversion(ts.units, self.units, bolo, wiringmap=self.wiringmap, hkmap=self.hkmap, system=self.system, tf=tf) except KeyError: if not self.skiperrors: raise newts[bolo] = ts continue self.convfactors[ts.units][bolo] = convfactor # And cache it # Convert timestream and store results convts = core.G3Timestream(ts) if core.G3TimestreamUnits.Resistance in [self.units, ts.units] and self.units != ts.units: convts = 1. / convts convts.units = self.units convts *= convfactor if convts.units != core.G3TimestreamUnits.Counts: convts.SetFLACCompression(False) newts[bolo] = convts frame[self.output] = newts
def subtract_dc_offset(frame, ts_key_in='RawTimestreams_I', ts_key_out='DCSubtractTimestreams'): if frame.type == core.G3FrameType.Scan and \ ts_key_in in frame.keys(): if ts_key_out not in frame.keys(): frame[ts_key_out] = core.G3TimestreamMap() for bolo in frame[ts_key_in].keys(): frame[ts_key_out][bolo] = frame[ts_key_in][bolo] - np.round( np.mean(frame[ts_key_in][bolo]))
def __call__(self, f): if f.type == core.G3FrameType.Scan: self.glitch_map = core.G3TimestreamMap() super().__call__(f) if f.type == core.G3FrameType.Scan: self.glitch_map.start = f[self.input].start self.glitch_map.stop = f[self.input].stop f[self.info] = self.glitch_map
def start_stream(self, session, params=None): if params is None: params = {} delay = params.get('delay', 1) ts_len = params.get('ts_len', 100) # Writes status frame f = core.G3Frame(core.G3FrameType.Housekeeping) f['session_id'] = 0 f['start_time'] = time.time() self.writer.Process(f) self.is_streaming = True frame_num = 0 while self.is_streaming: f = core.G3Frame(core.G3FrameType.Scan) t1 = time.time() t0 = t1 - delay ts = np.arange(t0, t1, 1 / self.freq) f['session_id'] = 0 f['frame_num'] = frame_num f['data'] = core.G3TimestreamMap() for k, c in self.channels.items(): fparams = copy.copy(c) bg = np.random.normal(0, fparams.get('stdev', 0), len(ts)) if fparams['type'] == 'const': xs = bg + fparams['val'] elif fparams['type'] in ['lin', 'linear']: xs = bg + ts * fparams['slope'] + fparams.get('offset', 0) # Wraps from -pi to pi xs = np.mod(xs + np.pi, 2 * np.pi) - np.pi f['data'][k] = core.G3Timestream(xs) f['data'][k].start = core.G3Time(t0 * core.G3Units.sec) f['data'][k].stop = core.G3Time(t1 * core.G3Units.sec) self.log.info("Writing G3 Frame") self.writer.Process(f) frame_num += 1 time.sleep(delay) print("Writing EndProcessingFrame") f = core.G3Frame(core.G3FrameType.EndProcessing) self.writer.Process(f) return True, "Finished streaming"
def start_stream(self, session, params=None): """ Task to stream fake detector data as G3Frames Args: frame_rate (float, optional): Frequency [Hz] at which G3Frames are sent over the network. Defaults to 1 frame pers sec. sample_rate (float, optional): Sample rate [Hz] for each channel. Defaults to 10 Hz. """ if params is None: params = {} frame_rate = params.get('frame_rate', 1.) sample_rate = params.get('sample_rate', 10.) f = core.G3Frame(core.G3FrameType.Observation) f['session_id'] = 0 f['start_time'] = time.time() self.writer.Process(f) frame_num = 0 self.is_streaming = True while self.is_streaming: frame_start = time.time() time.sleep(1. / frame_rate) frame_stop = time.time() times = np.arange(frame_start, frame_stop, 1. / sample_rate) f = core.G3Frame(core.G3FrameType.Scan) f['session_id'] = 0 f['frame_num'] = frame_num f['data'] = core.G3TimestreamMap() for i, chan in enumerate(self.channels): ts = core.G3Timestream([chan.read(t) for t in times]) ts.start = core.G3Time(frame_start * core.G3Units.sec) ts.stop = core.G3Time(frame_stop * core.G3Units.sec) f['data'][str(i)] = ts self.writer.Process(f) self.log.info("Writing frame...") frame_num += 1 return True, "Finished streaming"
def check_for_sim_keys(fr, valid_ids_key = 'valid_ids', out_tsm_key = 'valid_ids_tsm'): ''' CalculatePointing expects a TimestreamMap, so make sure there's one in the frame when doing mock-observations. valid_ids_key should already exist in simstub and point to a list. ''' if fr.type == core.G3FrameType.Scan: if not valid_ids_key in fr: raise KeyError( "%s not found in frame. "%valid_ids_key + "List of valid bolometer ids required for mock-observing.") tsm = core.G3TimestreamMap() for bid in fr[valid_ids_key]: tsm[bid]=core.G3Timestream() fr[out_tsm_key] = tsm
def __call__(self, f): if f.type == core.G3FrameType.Scan: if self.input not in f.keys() or type(f[self.input]) != core.G3TimestreamMap: raise ValueError("""Frame is a Scan but does not have a G3Timestream map named {}""".format(self.input)) processing = core.G3TimestreamMap() for k in f[self.input].keys(): processing[k] = core.G3Timestream( self.process(f[self.input][k], k) ) processing.start = f[self.input].start processing.stop = f[self.input].stop if self.input == self.output: f.pop(self.input) f[self.output] = processing
def AddBenchData(f): ''' Add the optical bench positions to the frame. ''' if f.type != core.G3FrameType.GcpSlow: return bench_axes = ['y1', 'y2', 'y3', 'x4', 'x5', 'z6'] benchcom = core.G3TimestreamMap() benchpos = core.G3TimestreamMap() benchzero = core.G3TimestreamMap() benchoff = core.G3TimestreamMap() bencherr = core.G3TimestreamMap() bench_info = core.G3TimestreamMap() for i, key in enumerate(bench_axes): # As of 2017-08-03, SCU time is not trustworthy # start = f['antenna0']['scu']['benchSampleTime'][0][0] # stop = f['antenna0']['scu']['benchSampleTime'][0][-1] # For now, do this bit of evil start = f['antenna0']['tracker']['utc'][0][0] stop = f['antenna0']['tracker']['utc'][0][-1] benchcom[key] = core.G3Timestream( f['antenna0']['scu']['benchExpected'][i]) benchcom[key].start = start benchcom[key].stop = stop benchpos[key] = core.G3Timestream( f['antenna0']['scu']['benchActual'][i]) benchpos[key].start = start benchpos[key].stop = stop benchzero[key] = core.G3Timestream( f['antenna0']['scu']['benchZeros'][i]) benchzero[key].start = start benchzero[key].stop = stop benchoff[key] = core.G3Timestream( f['antenna0']['scu']['benchOffsets'][i]) benchoff[key].start = start benchoff[key].stop = stop bencherr[key] = core.G3Timestream( f['antenna0']['scu']['benchErrors'][i]) bencherr[key].start = start bencherr[key].stop = stop info_items = [ 'benchFocus', 'benchDeadBand', 'benchAcquiredThreshold', 'benchPrimaryState', 'benchSecondaryState', 'benchFault', 'timeLocked' ] bench_info = core.G3TimestreamMap() for i, key in enumerate(info_items): start = f['antenna0']['tracker']['utc'][0][0] stop = f['antenna0']['tracker']['utc'][0][-1] bench_info[key] = core.G3Timestream(f['antenna0']['scu'][key][0]) bench_info[key].start = start bench_info[key].stop = stop f['BenchPosition'] = benchpos f['BenchCommandedPosition'] = benchcom f['BenchZeros'] = benchzero f['BenchOffsets'] = benchoff f['BenchErrors'] = bencherr f['BenchInfo'] = bench_info f['BenchSampleTime'] = f['antenna0']['scu']['benchSampleTime'][0]
def start_background_streamer(self, session, params=None): """start_background_streamer(params=None) Process to run streaming process. A data stream is started automatically. It can be stopped and started by the start and stop tasks. Either way keep alive flow control frames are being sent. Parameters ---------- frame_rate : float, optional Frequency [Hz] at which G3Frames are sent over the network. Defaults to 1 frame pers sec. sample_rate : float, optional Sample rate [Hz] for each channel. Defaults to 10 Hz. """ if params is None: params = {} self.writer = core.G3NetworkSender(hostname=self.target_host, port=self.port) frame_rate = params.get('frame_rate', 1.) sample_rate = params.get('sample_rate', 10.) frame_num = 0 self.running_in_background = True # Control flags FIFO stack to keep Writer single threaded self.flags = deque([FlowControl.START]) while self.running_in_background: # Send START frame if next(iter(self.flags), None) is FlowControl.START: self._set_stream_on() # sends start flowcontrol self.is_streaming = True self.flags.popleft() print("stream running in background") self.log.debug("control flags: {f}", f=self.flags) # Send keep alive flow control frame f = core.G3Frame(core.G3FrameType.none) f['sostream_flowcontrol'] = FlowControl.ALIVE.value self.writer.Process(f) if self.is_streaming: frame_start = time.time() time.sleep(1. / frame_rate) frame_stop = time.time() times = np.arange(frame_start, frame_stop, 1. / sample_rate) f = core.G3Frame(core.G3FrameType.Scan) f['session_id'] = 0 f['frame_num'] = frame_num f['sostream_id'] = self.stream_id f['data'] = core.G3TimestreamMap() for i, chan in enumerate(self.channels): ts = core.G3Timestream([chan.read() for t in times]) ts.start = core.G3Time(frame_start * core.G3Units.sec) ts.stop = core.G3Time(frame_stop * core.G3Units.sec) f['data'][f"r{i:04}"] = ts self.writer.Process(f) self.log.info("Writing frame...") frame_num += 1 # Send END frame if next(iter(self.flags), None) is FlowControl.END: self._send_end_flowcontrol_frame() self._send_cleanse_flowcontrol_frame() self.is_streaming = False self.flags.popleft() else: # Don't send keep alive frames too quickly time.sleep(1) # Shutdown streamer if next(iter(self.flags), None) is SHUTDOWN: self.running_in_background = False self.flags.popleft() # Teardown writer self.writer.Close() self.writer = None return True, "Finished streaming"
def __call__(self, f): if f.type == FT.Calibration and f['cal_type'] == 'focal_plane': self.focal_plane = f if f.type != FT.Scan: return [f] if f.type == FT.EndProcessing: flush = True else: flush = False self.frame_buffer.append(f) self.raw_buffer.append(f[self.enc_name]) # Figure out what frames we're able to process, given info we have. frames_out = [] # Work in units of seconds. raw_t0, raw_t1 = self.raw_buffer[0].t[0], self.raw_buffer[-1].t[-1] # Process any frame that ends before raw_t1. frame_index = 0 while len(self.frame_buffer) > 0: f = self.frame_buffer[0] if not flush and (f['signal'].stop.time / core.G3Units.sec > raw_t1): break sig = f[self.signal_name] # f['signal'] frame_t0 = sig.start.time / core.G3Units.sec frame_t1 = sig.stop.time / core.G3Units.sec tick_rate = sig.sample_rate / core.G3Units.Hz # Figure out what range of samples we will be able to set. start_index = np.ceil((raw_t0 - frame_t0) * tick_rate) end_index = np.floor((raw_t1 - frame_t0) * tick_rate) + 1 start_index = max(0, int(start_index)) end_index = min(int(end_index), sig.n_samples) if end_index != sig.n_samples and not flush: # Buffer. break # Otherwise, do the interpolations... frames_out.append(self.frame_buffer.pop(0)) t_raw = np.hstack([r.t for r in self.raw_buffer]) t_int = frame_t0 + np.arange(start_index, end_index) / tick_rate boresight = core.G3TimestreamMap() vs = {} for k in ['az', 'el', 'corot']: interp = spline1d( t_raw, np.hstack([r.data[k] for r in self.raw_buffer])) v = np.empty(sig.n_samples) v[:start_index] = np.nan v[start_index:end_index] = interp(t_int) v[end_index:] = np.nan vs[k] = v boresight[k] = core.G3Timestream(vs[k]) boresight.start = sig.start boresight.stop = sig.stop f[self.boresight_name] = boresight # f['boresight'] # Compute quaternion. q = ( # Sky <-- near (el, az=0) coords.q_euler(2, -vs['az'] * np.pi / 180) * # ... sky at az=0 <-- near (el=0,az=0) coords.q_euler(1, -vs['el'] * np.pi / 180) * # ... (1,-xi,eta) <-- (-eta,-xi,1) coords.q_euler(1, np.pi / 2) * # ... (-eta,-xi,1) <-- (eta,xi,1) coords.q_euler(2, np.pi)) # Note that there's no "TimestreamQuat" class. So no timestamps. f[self.boresight_name + '_q'] = q # f['boresight_q'] # Discard raw data we're not using any more. Out of caution, # keep one more frame than we have buffered. while len(self.raw_buffer) - len(self.frame_buffer) > 2: self.raw_buffer.pop(0) return frames_out
def cache_to_frames(tod, start_frame, n_frames, frame_offsets, frame_sizes, common=None, detector_fields=None, flag_fields=None, detector_map="detectors", flag_map="flags", units=None): """Gather all data from the distributed cache for a single frame. Args: tod (toast.TOD): instance of a TOD class. start_frame (int): the first frame index. n_frames (int): the number of frames. frame_offsets (list): list of the first samples of all frames. frame_sizes (list): list of the number of samples in each frame. common (tuple): (cache name, G3 type, frame name) of each common field. detector_fields (tuple): (cache name, frame name) of each detector field. flag_fields (tuple): (cache name, frame name) of each flag field. detector_map (str): the name of the frame timestream map. flag_map (str): then name of the frame flag map. units: G3 units of the detector data. """ # Local sample range local_first = tod.local_samples[0] nlocal = tod.local_samples[1] # The process grid detranks, sampranks = tod.grid_size rankdet, ranksamp = tod.grid_ranks # Helper function: # For a given timestream, the gather is done across the # process row which contains the specific detector, or across # the first process row for common telescope data. def gather_field(prow, fld, indx, cacheoff, ncache): gproc = 0 gdata = None # We are going to allreduce this later, so that every process # knows the dimensions of the field. allnnz = 0 if rankdet == prow: #print(" proc {} doing gather of {}".format(tod.mpicomm.rank, fld), flush=True) # This process is in the process row that has this field, # participate in the gather operation. pdata = None # Find the data type and shape from the cache object mtype = None ref = tod.cache.reference(fld) nnz = 1 if (len(ref.shape) > 1) and (ref.shape[1] > 0): nnz = ref.shape[1] if ref.dtype == np.dtype(np.float64): mtype = MPI.DOUBLE elif ref.dtype == np.dtype(np.int64): mtype = MPI.INT64_T elif ref.dtype == np.dtype(np.int32): mtype = MPI.INT32_T elif ref.dtype == np.dtype(np.uint8): mtype = MPI.UINT8_T else: msg = "Cannot gather cache field {} of type {}"\ .format(fld, ref.dtype) raise RuntimeError(msg) #print("field {}: proc {} has nnz = {}".format(fld, tod.mpicomm.rank, nnz), flush=True) pz = 0 if cacheoff is not None: pdata = ref.flatten()[nnz * cacheoff:nnz * (cacheoff + ncache)] pz = nnz * ncache else: pdata = np.zeros(0, dtype=ref.dtype) psizes = tod.grid_comm_row.gather(pz, root=0) disp = None totsize = None if ranksamp == 0: #print("Gathering field {} with type {}".format(fld, mtype), flush=True) # We are the process collecting the gathered data. gproc = tod.mpicomm.rank allnnz = nnz # Compute the displacements into the receive buffer. disp = [0] for ps in psizes[:-1]: last = disp[-1] disp.append(last + ps) totsize = np.sum(psizes) # allocate receive buffer gdata = np.zeros(totsize, dtype=ref.dtype) #print("Gatherv psizes = {}, disp = {}".format(psizes, disp), flush=True) #print("field {}: proc {} start Gatherv".format(fld, tod.mpicomm.rank), flush=True) tod.grid_comm_row.Gatherv(pdata, [gdata, psizes, disp, mtype], root=0) #print("field {}: proc {} finish Gatherv".format(fld, tod.mpicomm.rank), flush=True) del disp del psizes del pdata del ref # Now send this data to the root process of the whole communicator. # Only one process (the first one in process row "prow") has data # to send. # Create a unique message tag mtag = 10 * indx #print(" proc {} hit allreduce of gproc".format(tod.mpicomm.rank), flush=True) # All processes find out which one did the gather gproc = tod.mpicomm.allreduce(gproc, MPI.SUM) # All processes find out the field dimensions allnnz = tod.mpicomm.allreduce(allnnz, MPI.SUM) #print(" proc {} for field {}, gproc = {}".format(tod.mpicomm.rank, fld, gproc), flush=True) #print("field {}: proc {}, gatherproc = {}, allnnz = {}".format(fld, tod.mpicomm.rank, gproc, allnnz), flush=True) rdata = None if gproc == 0: if gdata is not None: if allnnz == 1: rdata = gdata else: rdata = gdata.reshape((-1, allnnz)) else: # Data not yet on rank 0 if tod.mpicomm.rank == 0: # Receive data from the first process in this row #print(" proc {} for field {}, recv type".format(tod.mpicomm.rank, fld), flush=True) rtype = tod.mpicomm.recv(source=gproc, tag=(mtag + 1)) #print(" proc {} for field {}, recv size".format(tod.mpicomm.rank, fld), flush=True) rsize = tod.mpicomm.recv(source=gproc, tag=(mtag + 2)) #print(" proc {} for field {}, recv data".format(tod.mpicomm.rank, fld), flush=True) rdata = np.zeros(rsize, dtype=np.dtype(rtype)) tod.mpicomm.Recv(rdata, source=gproc, tag=mtag) # Reshape if needed if allnnz > 1: rdata = rdata.reshape((-1, allnnz)) elif (tod.mpicomm.rank == gproc): # Send our data #print(" proc {} for field {}, send {} samples of {}".format(tod.mpicomm.rank, fld, len(gdata), gdata.dtype.char), flush=True) #print(" proc {} for field {}, send type with tag = {}".format(tod.mpicomm.rank, fld, mtag+1), flush=True) tod.mpicomm.send(gdata.dtype.char, dest=0, tag=(mtag + 1)) #print(" proc {} for field {}, send size with tag = {}".format(tod.mpicomm.rank, fld, mtag+2), flush=True) tod.mpicomm.send(len(gdata), dest=0, tag=(mtag + 2)) #print(" proc {} for field {}, send data with tag {}".format(tod.mpicomm.rank, fld, mtag), flush=True) tod.mpicomm.Send(gdata, 0, tag=mtag) return rdata # For efficiency, we are going to gather the data for all frames at once. # Then we will split those up when doing the write. # Frame offsets relative to the memory buffers we are gathering fdataoff = [0] for f in frame_sizes[:-1]: last = fdataoff[-1] fdataoff.append(last + f) # The list of frames- only on the root process. fdata = None if tod.mpicomm.rank == 0: fdata = [c3g.G3Frame(c3g.G3FrameType.Scan) for f in range(n_frames)] else: fdata = [None for f in range(n_frames)] # Compute the overlap of all frames with the local process. We want to # to find the full sample range that this process overlaps the total set # of frames. cacheoff = None ncache = 0 for f in range(n_frames): # Compute overlap of the frame with the local samples. fcacheoff, froff, nfr = local_frame_indices(local_first, nlocal, frame_offsets[f], frame_sizes[f]) #print("proc {}: frame {} has cache off {}, fr off {}, nfr {}".format(tod.mpicomm.rank, f, fcacheoff, froff, nfr), flush=True) if fcacheoff is not None: if cacheoff is None: cacheoff = fcacheoff ncache = nfr else: ncache += nfr #print("proc {}: cache off now {}, ncache now {}".format(tod.mpicomm.rank, cacheoff, ncache), flush=True) # Now gather the full sample data one field at a time. The root process # splits up the results into frames. # First gather common fields from the first row of the process grid. for findx, (cachefield, g3t, framefield) in enumerate(common): #print("proc {} entering gather_field(0, {}, {}, {}, {})".format(tod.mpicomm.rank, cachefield, findx, cacheoff, ncache), flush=True) data = gather_field(0, cachefield, findx, cacheoff, ncache) if tod.mpicomm.rank == 0: #print("Casting field {} to type {}".format(field, g3t), flush=True) if g3t == c3g.G3VectorTime: # Special case for time values stored as int64_t, but # wrapped in a class. for f in range(n_frames): dataoff = fdataoff[f] ndata = frame_sizes[f] g3times = list() for t in range(ndata): g3times.append(c3g.G3Time(data[dataoff + t])) fdata[f][framefield] = c3g.G3VectorTime(g3times) del g3times else: # The bindings of G3Vector seem to only work with # lists. This is probably horribly inefficient. for f in range(n_frames): dataoff = fdataoff[f] ndata = frame_sizes[f] if len(data.shape) == 1: fdata[f][framefield] = \ g3t(data[dataoff:dataoff+ndata].tolist()) else: # We have a 2D quantity fdata[f][framefield] = \ g3t(data[dataoff:dataoff+ndata,:].flatten().tolist()) del data # Wait for everyone to catch up... tod.mpicomm.barrier() # For each detector field, processes which have the detector # in their local_dets should be in the same process row. # We do the gather over just this process row. if (detector_fields is not None) or (flag_fields is not None): dpats = {d: re.compile(".*{}.*".format(d)) for d in tod.local_dets} detmaps = None if detector_fields is not None: if tod.mpicomm.rank == 0: detmaps = [c3g.G3TimestreamMap() for f in range(n_frames)] for dindx, (cachefield, framefield) in enumerate(detector_fields): pc = -1 for det, pat in dpats.items(): if pat.match(cachefield) is not None: #print("proc {} has field {}".format(tod.mpicomm.rank, field), flush=True) pc = rankdet break # As a sanity check, verify that every process which # has this field is in the same process row. rowcheck = tod.mpicomm.gather(pc, root=0) prow = 0 if tod.mpicomm.rank == 0: rc = np.array([x for x in rowcheck if (x >= 0)], dtype=np.int32) #print(field, rc, flush=True) prow = np.max(rc) if np.min(rc) != prow: msg = "Processes with field {} are not in the "\ "same row\n".format(cachefield) sys.stderr.write(msg) tod.mpicomm.abort() # Every process finds out which process row is participating. prow = tod.mpicomm.bcast(prow, root=0) #print("proc {} got prow = {}".format(tod.mpicomm.rank, prow), flush=True) # Get the data on rank 0 data = gather_field(prow, cachefield, dindx, cacheoff, ncache) if tod.mpicomm.rank == 0: if units is None: # We do this conditional, since we can't use # G3TimestreamUnits.None in python ("None" is # interpreted as python None). for f in range(n_frames): dataoff = fdataoff[f] ndata = frame_sizes[f] detmaps[f][framefield] = \ c3g.G3Timestream(data[dataoff:dataoff+ndata]) else: for f in range(n_frames): dataoff = fdataoff[f] ndata = frame_sizes[f] detmaps[f][framefield] = \ c3g.G3Timestream(data[dataoff:dataoff+ndata], units) if tod.mpicomm.rank == 0: for f in range(n_frames): fdata[f][detector_map] = detmaps[f] flagmaps = None if flag_fields is not None: if tod.mpicomm.rank == 0: flagmaps = [c3g.G3MapVectorInt() for f in range(n_frames)] for dindx, (cachefield, framefield) in enumerate(flag_fields): pc = -1 for det, pat in dpats.items(): if pat.match(cachefield) is not None: pc = rankdet break # As a sanity check, verify that every process which # has this field is in the same process row. rowcheck = tod.mpicomm.gather(pc, root=0) prow = 0 if tod.mpicomm.rank == 0: rc = np.array([x for x in rowcheck if (x >= 0)], dtype=np.int32) prow = np.max(rc) if np.min(rc) != prow: msg = "Processes with field {} are not in the "\ "same row\n".format(cachefield) sys.stderr.write(msg) tod.mpicomm.abort() # Every process finds out which process row is participating. prow = tod.mpicomm.bcast(prow, root=0) # Get the data on rank 0 data = gather_field(prow, cachefield, dindx, cacheoff, ncache) if tod.mpicomm.rank == 0: # The bindings of G3Vector seem to only work with # lists... Also there is no vectormap for unsigned # char, so we have to use int... for f in range(n_frames): dataoff = fdataoff[f] ndata = frame_sizes[f] flagmaps[f][framefield] = \ c3g.G3VectorInt(\ data[dataoff:dataoff+ndata].astype(np.int32)\ .tolist()) if tod.mpicomm.rank == 0: for f in range(n_frames): fdata[f][flag_map] = flagmaps[f] return fdata
def tod_to_frames( tod, start_frame, n_frames, frame_offsets, frame_sizes, cache_signal=None, cache_flags=None, cache_common_flags=None, copy_common=None, copy_detector=None, mask_flag_common=255, mask_flag=255, units=None, dets=None, compress=False, ): """Gather all data from the distributed TOD cache for a set of frames. Args: tod (toast.TOD): instance of a TOD class. start_frame (int): the first frame index. n_frames (int): the number of frames. frame_offsets (array_like): list of the first samples of all frames. frame_sizes (list): list of the number of samples in each frame. cache_signal (str): if None, read signal from TOD. Otherwise use this cache prefix for the detector signal timestreams. cache_flags (str): if None read det flags from TOD. Otherwise use this cache prefix for the detector flag timestreams. cache_common_flags (str): if None, read common flags from TOD. Otherwise use this cache prefix. copy_common (tuple): (cache name, G3 type, frame name) of each extra common field to copy from cache. copy_detector (tuple): (cache name prefix, G3 type, G3 map type, frame name) of each distributed detector field (excluding the "signal") to copy from cache. mask_flag_common (int): Bitmask to apply to common flags. mask_flag (int): Bitmask to apply to per-detector flags. units: G3 units of the detector data. dets (list): List of detectors to include in the frame. If None, use all of the detectors in the TOD object. compress (bool or dict): If True or a dictionary of compression parameters, store the timestreams as FLAC-compressed, 24-bit integers instead of uncompressed doubles. Returns: (list): List of frames on rank zero. Other processes have a list of None values. """ comm = tod.mpicomm rank = 0 if comm is not None: rank = comm.rank comm_row = tod.grid_comm_row # Detector names if dets is None: detnames = tod.detectors else: detnames = [] use_dets = set(dets) for det in tod.detectors: if det in use_dets: detnames.append(det) # Local sample range local_first = tod.local_samples[0] nlocal = tod.local_samples[1] # The process grid detranks, sampranks = tod.grid_size rankdet, ranksamp = tod.grid_ranks def get_local_cache(prow, field, cacheoff, ncache): """Read a local slice of a cache field. """ mtype = None pdata = None nnz = 0 if rankdet == prow: ref = tod.cache.reference(field) nnz = 1 if (len(ref.shape) > 1) and (ref.shape[1] > 0): nnz = ref.shape[1] if comm is not None: if ref.dtype == np.dtype(np.float64): mtype = MPI.DOUBLE elif ref.dtype == np.dtype(np.int64): mtype = MPI.INT64_T elif ref.dtype == np.dtype(np.int32): mtype = MPI.INT32_T elif ref.dtype == np.dtype(np.uint8): mtype = MPI.UINT8_T else: msg = "Cannot use cache field {} of type {}"\ .format(field, ref.dtype) raise RuntimeError(msg) if cacheoff is not None: pdata = ref.flatten()[nnz * cacheoff:nnz * (cacheoff + ncache)] else: pdata = np.zeros(0, dtype=ref.dtype) return (pdata, nnz, mtype) def gather_field(prow, pdata, nnz, mpitype, cacheoff, ncache, tag): """Gather a single timestream buffer to the root process. """ is_none = pdata is None all_none = comm.allreduce(is_none, MPI.LAND) if all_none: # This situation arises at least when gathering HWP angle from LAT return None gdata = None # We are going to allreduce this later, so that every process # knows the dimensions of the field. gproc = 0 allnnz = 0 # Size of the local buffer pz = 0 if pdata is not None: pz = len(pdata) if rankdet == prow: psizes = None if comm_row is None: psizes = [pz] else: psizes = comm_row.gather(pz, root=0) disp = None totsize = None if ranksamp == 0: # We are the process collecting the gathered data. allnnz = nnz gproc = rank # Compute the displacements into the receive buffer. disp = [0] for ps in psizes[:-1]: last = disp[-1] disp.append(last + ps) totsize = np.sum(psizes) # allocate receive buffer gdata = np.zeros(totsize, dtype=pdata.dtype) if comm_row is None: pdata[:] = gdata else: comm_row.Gatherv(pdata, [gdata, psizes, disp, mpitype], root=0) del disp del psizes # Now send this data to the root process of the whole communicator. # Only one process (the first one in process row "prow") has data # to send. if comm is not None: # All processes find out which one did the gather gproc = comm.allreduce(gproc, MPI.SUM) # All processes find out the field dimensions allnnz = comm.allreduce(allnnz, MPI.SUM) mtag = 10 * tag rdata = None if gproc == 0: if gdata is not None: if allnnz == 1: rdata = gdata else: rdata = gdata.reshape((-1, allnnz)) else: # Data not yet on rank 0 if rank == 0: # Receive data from the first process in this row rtype = comm.recv(source=gproc, tag=(mtag + 1)) rsize = comm.recv(source=gproc, tag=(mtag + 2)) rdata = np.zeros(rsize, dtype=np.dtype(rtype)) comm.Recv(rdata, source=gproc, tag=mtag) # Reshape if needed if allnnz > 1: rdata = rdata.reshape((-1, allnnz)) elif (rank == gproc): # Send our data comm.send(gdata.dtype.char, dest=0, tag=(mtag + 1)) comm.send(len(gdata), dest=0, tag=(mtag + 2)) comm.Send(gdata, 0, tag=mtag) return rdata # For efficiency, we are going to gather the data for all frames at once. # Then we will split those up when doing the write. # Frame offsets relative to the memory buffers we are gathering fdataoff = [0] for f in frame_sizes[:-1]: last = fdataoff[-1] fdataoff.append(last + f) # The list of frames- only on the root process. fdata = None if rank == 0: fdata = [ core3g.G3Frame(core3g.G3FrameType.Scan) for f in range(n_frames) ] else: fdata = [None for f in range(n_frames)] def split_field(data, g3t, framefield, mapfield=None, g3units=units, times=None): """Split a gathered data buffer into frames- only on root process. """ if data is None: return if g3t == core3g.G3VectorTime: # Special case for time values stored as int64_t, but # wrapped in a class. for f in range(n_frames): dataoff = fdataoff[f] ndata = frame_sizes[f] g3times = list() for t in range(ndata): g3times.append(core3g.G3Time(data[dataoff + t])) if mapfield is None: fdata[f][framefield] = core3g.G3VectorTime(g3times) else: fdata[f][framefield][mapfield] = \ core3g.G3VectorTime(g3times) del g3times elif g3t == so3g.IntervalsInt: # Flag vector is written as a simple boolean. for f in range(n_frames): dataoff = fdataoff[f] ndata = frame_sizes[f] # Extract flag vector (0 or 1) for this frame frame_flags = (data[dataoff:dataoff + ndata] != 0).astype(int) # Convert bit 0 to an IntervalsInt. ival = so3g.IntervalsInt.from_mask(frame_flags, 1)[0] if mapfield is None: fdata[f][framefield] = ival else: fdata[f][framefield][mapfield] = ival elif g3t == core3g.G3Timestream: if times is None: raise RuntimeError( "You must provide the time stamp vector with a " "Timestream object") for f in range(n_frames): dataoff = fdataoff[f] ndata = frame_sizes[f] timeslice = times[cacheoff + dataoff:cacheoff + dataoff + ndata] tstart = timeslice[0] * 1e8 tstop = timeslice[-1] * 1e8 if mapfield is None: if g3units is None: fdata[f][framefield] = \ g3t(data[dataoff : dataoff + ndata]) else: fdata[f][framefield] = \ g3t(data[dataoff : dataoff + ndata], g3units) fdata[f][framefield].start = core3g.G3Time(tstart) fdata[f][framefield].stop = core3g.G3Time(tstop) else: # Individual detector data. The only fields that # we (optionally) compress. if g3units is None: tstream = g3t(data[dataoff:dataoff + ndata]) else: tstream = g3t(data[dataoff:dataoff + ndata], g3units) if compress and "compressor_gain_" + framefield in fdata[f]: (tstream, gain, offset) = recode_timestream(tstream, compress) fdata[f]["compressor_gain_" + framefield][mapfield] = gain fdata[f]["compressor_offset_" + framefield][mapfield] = offset fdata[f][framefield][mapfield] = tstream fdata[f][framefield][mapfield].start = core3g.G3Time( tstart) fdata[f][framefield][mapfield].stop = core3g.G3Time(tstop) else: # The bindings of G3Vector seem to only work with # lists. This is probably horribly inefficient. for f in range(n_frames): dataoff = fdataoff[f] ndata = frame_sizes[f] if len(data.shape) == 1: fdata[f][framefield] = \ g3t(data[dataoff : dataoff + ndata].tolist()) else: # We have a 2D quantity fdata[f][framefield] = \ g3t(data[dataoff : dataoff + ndata, :].flatten() .tolist()) return # Compute the overlap of all frames with the local process. We want to # to find the full sample range that this process overlaps the total set # of frames. cacheoff = None ncache = 0 for f in range(n_frames): # Compute overlap of the frame with the local samples. fcacheoff, froff, nfr = s3utils.local_frame_indices( local_first, nlocal, frame_offsets[f], frame_sizes[f]) if fcacheoff is not None: if cacheoff is None: cacheoff = fcacheoff ncache = nfr else: ncache += nfr # Now gather the full sample data one field at a time. The root process # splits up the results into frames. # First collect boresight data. In addition to quaternions for the Az/El # pointing, we convert this back into angles that follow the specs # for telescope pointing. times = None if rankdet == 0: times = tod.local_times() if comm is not None: times = gather_field(0, times, 1, MPI.DOUBLE, cacheoff, ncache, 0) bore = None if rankdet == 0: bore = tod.read_boresight(local_start=cacheoff, n=ncache).flatten() if comm is not None: bore = gather_field(0, bore, 4, MPI.DOUBLE, cacheoff, ncache, 0) if rank == 0: split_field(bore.reshape(-1, 4), core3g.G3VectorDouble, "boresight_radec") bore = None if rankdet == 0: bore = tod.read_boresight_azel(local_start=cacheoff, n=ncache).flatten() if comm is not None: bore = gather_field(0, bore, 4, MPI.DOUBLE, cacheoff, ncache, 1) if rank == 0: split_field(bore.reshape(-1, 4), core3g.G3VectorDouble, "boresight_azel") corotator_angle = None if rankdet == 0: cache_name = "corotator_angle_deg" if tod.cache.exists(cache_name): corotator_angle = np.radians(tod.cache.reference(cache_name)) if comm is not None: corotator_angle = gather_field(0, corotator_angle, 1, MPI.DOUBLE, cacheoff, ncache, 0) if rank == 0: split_field(corotator_angle, core3g.G3VectorDouble, "corotator_angle", times=times) if rank == 0: for f in range(n_frames): fdata[f]["boresight"] = core3g.G3TimestreamMap() ang_theta, ang_phi, ang_psi = qa.to_angles(bore) # Astronomical convention for azimuth is opposite to spherical # coordinate phi. ang_az = -ang_phi ang_el = (np.pi / 2.0) - ang_theta ang_roll = ang_psi split_field(ang_az, core3g.G3Timestream, "boresight", "az", None, times=times) split_field(ang_el, core3g.G3Timestream, "boresight", "el", None, times=times) split_field(ang_roll, core3g.G3Timestream, "boresight", "roll", None, times=times) hwp_angle = None if rankdet == 0: hwp_angle = tod.local_hwp_angle() if comm is not None: hwp_angle = gather_field(0, hwp_angle, 1, MPI.DOUBLE, cacheoff, ncache, 0) if rank == 0: split_field(hwp_angle, core3g.G3VectorDouble, "hwp_angle", times=times) # Now the position and velocity information pos = None if rankdet == 0: pos = tod.read_position(local_start=cacheoff, n=ncache).flatten() if comm is not None: pos = gather_field(0, pos, 3, MPI.DOUBLE, cacheoff, ncache, 2) if rank == 0: split_field(pos.reshape(-1, 3), core3g.G3VectorDouble, "site_position") vel = None if rankdet == 0: vel = tod.read_velocity(local_start=cacheoff, n=ncache).flatten() if comm is not None: vel = gather_field(0, vel, 3, MPI.DOUBLE, cacheoff, ncache, 3) if rank == 0: split_field(vel.reshape(-1, 3), core3g.G3VectorDouble, "site_velocity") # Now handle the common flags- either from a cache object or from the # TOD methods cflags = None nnz = 1 if cache_common_flags is None: if rankdet == 0: cflags = np.array( tod.read_common_flags(local_start=cacheoff, n=ncache)) cflags &= mask_flag_common else: cflags, nnz, mtype = get_local_cache(0, cache_common_flags, cacheoff, ncache) if cflags is not None: cflags &= mask_flag_common if comm is not None: mtype = MPI.UINT8_T cflags = gather_field(0, cflags, nnz, mtype, cacheoff, ncache, 4) if rank == 0: split_field(cflags, so3g.IntervalsInt, "flags_common") # Any extra common fields if comm is not None: comm.barrier() if copy_common is not None: for cindx, (cname, g3typ, fname) in enumerate(copy_common): cdata, nnz, mtype = get_local_cache(0, cname, cacheoff, ncache) cdata = gather_field(0, cdata, nnz, mtype, cacheoff, ncache, cindx) if rank == 0: split_field(cdata, g3typ, fname) # Now read all per-detector quantities. # For each detector field, processes which have the detector # in their local_dets should be in the same process row. if rank == 0: for f in range(n_frames): fdata[f]["signal"] = core3g.G3TimestreamMap() if compress: fdata[f]["compressor_gain_signal"] = core3g.G3MapDouble() fdata[f]["compressor_offset_signal"] = core3g.G3MapDouble() fdata[f]["flags"] = so3g.MapIntervalsInt() if copy_detector is not None: for cname, g3typ, g3maptyp, fnm in copy_detector: fdata[f][fnm] = g3maptyp() if compress: fdata[f]["compressor_gain_" + fnm] = core3g.G3MapDouble() fdata[f]["compressor_offset_" + fnm] = core3g.G3MapDouble() for dindx, dname in enumerate(detnames): drow = -1 # Demodulation may have synthesized new detector names dnames = [] for x in tod.local_dets: if x.endswith(dname): dnames.append(x) if dnames: drow = rankdet # As a sanity check, verify that every process which # has this detector is in the same process row. rowcheck = None if comm is None: rowcheck = [drow] else: rowcheck = comm.gather(drow, root=0) prow = 0 if rank == 0: rc = np.array([x for x in rowcheck if (x >= 0)], dtype=np.int32) prow = np.max(rc) if np.min(rc) != prow: msg = "Processes with detector {} are not in the "\ "same row of the process grid\n".format(dname) sys.stderr.write(msg) if comm is not None: comm.abort() # Every process finds out which process row is participating. if comm is not None: prow = comm.bcast(prow, root=0) all_dnames = comm.allgather(dnames) dnames = list(set(np.concatenate(all_dnames).flat)) # "signal" for dname in dnames: detdata = None nnz = 1 if cache_signal is None: if rankdet == prow: detdata = tod.local_signal(dname)[cacheoff:cacheoff + ncache] else: cache_det = "{}_{}".format(cache_signal, dname) detdata, nnz, mtype = get_local_cache(prow, cache_det, cacheoff, ncache) if comm is not None: mtype = MPI.DOUBLE detdata = gather_field(prow, detdata, nnz, mtype, cacheoff, ncache, dindx) if rank == 0: split_field(detdata, core3g.G3Timestream, "signal", mapfield=dname, times=times) # "flags" for dname in dnames: detdata = None nnz = 1 if cache_flags is None: if rankdet == prow: detdata = tod.local_flags(dname)[cacheoff:cacheoff + ncache] detdata &= mask_flag else: cache_det = "{}_{}".format(cache_flags, dname) detdata, nnz, mtype = get_local_cache(prow, cache_det, cacheoff, ncache) if detdata is not None: detdata &= mask_flag if comm is not None: mtype = MPI.UINT8_T detdata = gather_field(prow, detdata, nnz, mtype, cacheoff, ncache, dindx) if rank == 0: split_field(detdata, so3g.IntervalsInt, "flags", mapfield=dname) # Now copy any additional fields. for dname in dnames: if copy_detector is not None: for cname, g3typ, g3maptyp, fnm in copy_detector: cache_det = "{}_{}".format(cname, dname) detdata, nnz, mtype = get_local_cache( prow, cache_det, cacheoff, ncache) if comm is not None: detdata = gather_field(prow, detdata, nnz, mtype, cacheoff, ncache, dindx) if rank == 0: split_field(detdata, g3typ, fnm, mapfield=dname, times=times) return fdata
from spt3g import core then = core.G3Time.Now() now = core.G3Time(then.time + 3 * core.G3Units.second) data = numpy.zeros(200) ts = core.G3Timestream(data) ts.start = then ts.stop = now assert (numpy.abs(ts.sample_rate / core.G3Units.Hz - 66.33333) < 1e-5) times = ts.times() assert (ts.times()[0] == ts.start) print(ts.times()[-1].time, ts.stop.time) assert (numpy.abs(ts.times()[-1].time - ts.stop.time) < 2 ) # Max 1 tick roundoff error tsm = core.G3TimestreamMap() tsm['Test'] = ts assert (numpy.abs(tsm.sample_rate / core.G3Units.Hz - 66.33333) < 1e-5) times = tsm.times() assert (tsm.times()[0] == tsm.start) assert (numpy.abs(tsm.times()[-1].time - tsm.stop.time) < 2 ) # Max 1 tick roundoff error
def tod_to_frames(tod, start_frame, n_frames, frame_offsets, frame_sizes, cache_signal=None, cache_flags=None, cache_common_flags=None, copy_common=None, copy_detector=None, mask_flag_common=255, mask_flag=255, units=None): """Gather all data from the distributed TOD cache for a set of frames. Args: tod (toast.TOD): instance of a TOD class. start_frame (int): the first frame index. n_frames (int): the number of frames. frame_offsets (array_like): list of the first samples of all frames. frame_sizes (list): list of the number of samples in each frame. cache_signal (str): if None, read signal from TOD. Otherwise use this cache prefix for the detector signal timestreams. cache_flags (str): if None read det flags from TOD. Otherwise use this cache prefix for the detector flag timestreams. cache_common_flags (str): if None, read common flags from TOD. Otherwise use this cache prefix. copy_common (tuple): (cache name, G3 type, frame name) of each extra common field to copy from cache. copy_detector (tuple): (cache name prefix, G3 type, G3 map type, frame name) of each distributed detector field (excluding the "signal") to copy from cache. mask_flag_common (int): Bitmask to apply to common flags. mask_flag (int): Bitmask to apply to per-detector flags. units: G3 units of the detector data. Returns: (list): List of frames on rank zero. Other processes have a list of None values. """ # Detector names detnames = tod.detectors # Local sample range local_first = tod.local_samples[0] nlocal = tod.local_samples[1] # The process grid detranks, sampranks = tod.grid_size rankdet, ranksamp = tod.grid_ranks def get_local_cache(prow, fld, cacheoff, ncache): """Read a local slice of a cache field. """ mtype = None pdata = None if rankdet == prow: ref = tod.cache.reference(fld) nnz = 1 if (len(ref.shape) > 1) and (ref.shape[1] > 0): nnz = ref.shape[1] if ref.dtype == np.dtype(np.float64): mtype = MPI.DOUBLE elif ref.dtype == np.dtype(np.int64): mtype = MPI.INT64_T elif ref.dtype == np.dtype(np.int32): mtype = MPI.INT32_T elif ref.dtype == np.dtype(np.uint8): mtype = MPI.UINT8_T else: msg = "Cannot use cache field {} of type {}"\ .format(fld, ref.dtype) raise RuntimeError(msg) if cacheoff is not None: pdata = ref.flatten()[nnz * cacheoff:nnz * (cacheoff + ncache)] else: pdata = np.zeros(0, dtype=ref.dtype) return (pdata, nnz, mtype) def gather_field(prow, pdata, nnz, mpitype, cacheoff, ncache, tag): """Gather a single timestream buffer to the root process. """ gdata = None # We are going to allreduce this later, so that every process # knows the dimensions of the field. gproc = 0 allnnz = 0 # Size of the local buffer pz = len(pdata) if rankdet == prow: psizes = tod.grid_comm_row.gather(pz, root=0) disp = None totsize = None if ranksamp == 0: # We are the process collecting the gathered data. allnnz = nnz gproc = tod.mpicomm.rank # Compute the displacements into the receive buffer. disp = [0] for ps in psizes[:-1]: last = disp[-1] disp.append(last + ps) totsize = np.sum(psizes) # allocate receive buffer gdata = np.zeros(totsize, dtype=pdata.dtype) tod.grid_comm_row.Gatherv(pdata, [gdata, psizes, disp, mpitype], root=0) del disp del psizes # Now send this data to the root process of the whole communicator. # Only one process (the first one in process row "prow") has data # to send. # All processes find out which one did the gather gproc = tod.mpicomm.allreduce(gproc, MPI.SUM) # All processes find out the field dimensions allnnz = tod.mpicomm.allreduce(allnnz, MPI.SUM) mtag = 10 * tag rdata = None if gproc == 0: if gdata is not None: if allnnz == 1: rdata = gdata else: rdata = gdata.reshape((-1, allnnz)) else: # Data not yet on rank 0 if tod.mpicomm.rank == 0: # Receive data from the first process in this row rtype = tod.mpicomm.recv(source=gproc, tag=(mtag + 1)) rsize = tod.mpicomm.recv(source=gproc, tag=(mtag + 2)) rdata = np.zeros(rsize, dtype=np.dtype(rtype)) tod.mpicomm.Recv(rdata, source=gproc, tag=mtag) # Reshape if needed if allnnz > 1: rdata = rdata.reshape((-1, allnnz)) elif (tod.mpicomm.rank == gproc): # Send our data tod.mpicomm.send(gdata.dtype.char, dest=0, tag=(mtag + 1)) tod.mpicomm.send(len(gdata), dest=0, tag=(mtag + 2)) tod.mpicomm.Send(gdata, 0, tag=mtag) return rdata # For efficiency, we are going to gather the data for all frames at once. # Then we will split those up when doing the write. # Frame offsets relative to the memory buffers we are gathering fdataoff = [0] for f in frame_sizes[:-1]: last = fdataoff[-1] fdataoff.append(last + f) # The list of frames- only on the root process. fdata = None if tod.mpicomm.rank == 0: fdata = [ core3g.G3Frame(core3g.G3FrameType.Scan) for f in range(n_frames) ] else: fdata = [None for f in range(n_frames)] def flags_to_intervals(flgs): """Convert a flag vector to an interval list. """ groups = [[i for i, value in it] for key, it in itertools.groupby(enumerate(flgs), key=operator.itemgetter(1)) if key != 0] chunks = list() for grp in groups: chunks.append([grp[0], grp[-1]]) return chunks def split_field(data, g3t, framefield, mapfield=None, g3units=units): """Split a gathered data buffer into frames. """ if tod.mpicomm.rank == 0: if g3t == core3g.G3VectorTime: # Special case for time values stored as int64_t, but # wrapped in a class. for f in range(n_frames): dataoff = fdataoff[f] ndata = frame_sizes[f] g3times = list() for t in range(ndata): g3times.append(core3g.G3Time(data[dataoff + t])) if mapfield is None: fdata[f][framefield] = core3g.G3VectorTime(g3times) else: fdata[f][framefield][mapfield] = \ core3g.G3VectorTime(g3times) del g3times elif g3t == so3g.IntervalsInt: # This means that the data is actually flags # and we should convert it into a list of intervals. fint = flags_to_intervals(data) for f in range(n_frames): dataoff = fdataoff[f] ndata = frame_sizes[f] datalast = dataoff + ndata chunks = list() idomain = (0, ndata - 1) for intr in fint: # Interval sample ranges are defined relative to the # frame itself. cfirst = None clast = None if (intr[0] < datalast) and (intr[1] >= dataoff): # there is some overlap... if intr[0] < dataoff: cfirst = 0 else: cfirst = intr[0] - dataoff if intr[1] >= datalast: clast = ndata - 1 else: clast = intr[1] - dataoff chunks.append([cfirst, clast]) if mapfield is None: if len(chunks) == 0: fdata[f][framefield] = \ so3g.IntervalsInt() else: fdata[f][framefield] = \ so3g.IntervalsInt.from_array( np.array(chunks, dtype=np.int64)) fdata[f][framefield].domain = idomain else: if len(chunks) == 0: fdata[f][framefield][mapfield] = \ so3g.IntervalsInt() else: fdata[f][framefield][mapfield] = \ so3g.IntervalsInt.from_array( np.array(chunks, dtype=np.int64)) fdata[f][framefield][mapfield].domain = idomain del fint elif g3t == core3g.G3Timestream: for f in range(n_frames): dataoff = fdataoff[f] ndata = frame_sizes[f] if mapfield is None: if g3units is None: fdata[f][framefield] = \ g3t(data[dataoff:dataoff+ndata]) else: fdata[f][framefield] = \ g3t(data[dataoff:dataoff+ndata], g3units) else: if g3units is None: fdata[f][framefield][mapfield] = \ g3t(data[dataoff:dataoff+ndata]) else: fdata[f][framefield][mapfield] = \ g3t(data[dataoff:dataoff+ndata], g3units) else: # The bindings of G3Vector seem to only work with # lists. This is probably horribly inefficient. for f in range(n_frames): dataoff = fdataoff[f] ndata = frame_sizes[f] if len(data.shape) == 1: fdata[f][framefield] = \ g3t(data[dataoff:dataoff+ndata].tolist()) else: # We have a 2D quantity fdata[f][framefield] = \ g3t(data[dataoff:dataoff+ndata, :].flatten() .tolist()) return # Compute the overlap of all frames with the local process. We want to # to find the full sample range that this process overlaps the total set # of frames. cacheoff = None ncache = 0 for f in range(n_frames): # Compute overlap of the frame with the local samples. fcacheoff, froff, nfr = s3utils.local_frame_indices( local_first, nlocal, frame_offsets[f], frame_sizes[f]) if fcacheoff is not None: if cacheoff is None: cacheoff = fcacheoff ncache = nfr else: ncache += nfr # Now gather the full sample data one field at a time. The root process # splits up the results into frames. # First collect boresight data. In addition to quaternions for the Az/El # pointing, we convert this back into angles that follow the specs # for telescope pointing. bore = None if rankdet == 0: bore = tod.read_boresight(local_start=cacheoff, n=ncache) bore = gather_field(0, bore.flatten(), 4, MPI.DOUBLE, cacheoff, ncache, 0) split_field(bore.reshape(-1, 4), core3g.G3VectorDouble, "qboresight_radec") bore = None if rankdet == 0: bore = tod.read_boresight_azel(local_start=cacheoff, n=ncache) bore = gather_field(0, bore.flatten(), 4, MPI.DOUBLE, cacheoff, ncache, 1) split_field(bore.reshape(-1, 4), core3g.G3VectorDouble, "qboresight_azel") if tod.mpicomm.rank == 0: for f in range(n_frames): fdata[f]["boresight"] = core3g.G3TimestreamMap() ang_theta, ang_phi, ang_psi = qa.to_angles(bore) ang_az = ang_phi ang_el = (np.pi / 2.0) - ang_theta ang_roll = ang_psi split_field(ang_az, core3g.G3Timestream, "boresight", "az", None) split_field(ang_el, core3g.G3Timestream, "boresight", "el", None) split_field(ang_roll, core3g.G3Timestream, "boresight", "roll", None) # Now the position and velocity information pos = None if rankdet == 0: pos = tod.read_position(local_start=cacheoff, n=ncache) pos = gather_field(0, pos.flatten(), 3, MPI.DOUBLE, cacheoff, ncache, 2) split_field(pos.reshape(-1, 3), core3g.G3VectorDouble, "site_position") vel = None if rankdet == 0: vel = tod.read_velocity(local_start=cacheoff, n=ncache) vel = gather_field(0, vel.flatten(), 3, MPI.DOUBLE, cacheoff, ncache, 3) split_field(vel.reshape(-1, 3), core3g.G3VectorDouble, "site_velocity") # Now handle the common flags- either from a cache object or from the # TOD methods cflags = None nnz = 1 mtype = MPI.UINT8_T if cache_common_flags is None: if rankdet == 0: cflags = tod.read_common_flags(local_start=cacheoff, n=ncache) cflags &= mask_flag_common else: cflags, nnz, mtype = get_local_cache(0, cache_common_flags, cacheoff, ncache) cflags &= mask_flag_common cflags = gather_field(0, cflags, nnz, mtype, cacheoff, ncache, 4) split_field(cflags, so3g.IntervalsInt, "flags_common") # Any extra common fields tod.mpicomm.barrier() if copy_common is not None: for cindx, (cname, g3typ, fname) in enumerate(copy_common): cdata, nnz, mtype = get_local_cache(0, cname, cacheoff, ncache) cdata = gather_field(0, cdata, nnz, mtype, cacheoff, ncache, cindx) split_field(cdata, g3typ, fname) # Now read all per-detector quantities. # For each detector field, processes which have the detector # in their local_dets should be in the same process row. if tod.mpicomm.rank == 0: for f in range(n_frames): fdata[f]["signal"] = core3g.G3TimestreamMap() fdata[f]["flags"] = so3g.MapIntervalsInt() if copy_detector is not None: for cname, g3typ, g3maptyp, fnm in copy_detector: fdata[f][fnm] = g3maptyp() for dindx, dname in enumerate(detnames): drow = -1 if dname in tod.local_dets: drow = rankdet # As a sanity check, verify that every process which # has this detector is in the same process row. rowcheck = tod.mpicomm.gather(drow, root=0) prow = 0 if tod.mpicomm.rank == 0: rc = np.array([x for x in rowcheck if (x >= 0)], dtype=np.int32) prow = np.max(rc) if np.min(rc) != prow: msg = "Processes with detector {} are not in the "\ "same row of the process grid\n".format(dname) sys.stderr.write(msg) tod.mpicomm.abort() # Every process finds out which process row is participating. prow = tod.mpicomm.bcast(prow, root=0) # "signal" detdata = None nnz = 1 mtype = MPI.DOUBLE if cache_signal is None: if rankdet == prow: detdata = tod.read(detector=dname, local_start=cacheoff, n=ncache) else: cache_det = "{}_{}".format(cache_signal, dname) detdata, nnz, mtype = get_local_cache(prow, cache_det, cacheoff, ncache) detdata = gather_field(prow, detdata, nnz, mtype, cacheoff, ncache, dindx) split_field(detdata, core3g.G3Timestream, "signal", mapfield=dname) # "flags" detdata = None nnz = 1 mtype = MPI.UINT8_T if cache_flags is None: if rankdet == prow: detdata = tod.read_flags(detector=dname, local_start=cacheoff, n=ncache) detdata &= mask_flag else: cache_det = "{}_{}".format(cache_flags, dname) detdata, nnz, mtype = get_local_cache(prow, cache_det, cacheoff, ncache) detdata &= mask_flag detdata = gather_field(prow, detdata, nnz, mtype, cacheoff, ncache, dindx) split_field(detdata, so3g.IntervalsInt, "flags", mapfield=dname) # Now copy any additional fields. if copy_detector is not None: for cname, g3typ, g3maptyp, fnm in copy_detector: cache_det = "{}_{}".format(cname, dname) detdata, nnz, mtype = get_local_cache(prow, cache_det, cacheoff, ncache) detdata = gather_field(prow, detdata, nnz, mtype, cacheoff, ncache, dindx) split_field(detdata, g3typ, fnm, mapfield=dname) return fdata