class QueryDataspec(task.MPILoggedTask): """Find the available files given a dataspec in the config file. Attributes ---------- instrument : str Name of the instrument. timerange : list List of time ranges as documented above. node_spoof : dict, optional Optional node spoof argument. """ instrument = config.Property(proptype=str) timerange = config.Property(proptype=list) node_spoof = config.Property(proptype=dict, default=_DEFAULT_NODE_SPOOF) def setup(self): """Fetch the files in the given dataspec. Returns ------- files : list List of files to load """ dspec = {"instrument": self.instrument, "timerange": self.timerange} # Add archive root if exists if self.node_spoof is not None: dspec["node_spoof"] = self.node_spoof files = files_from_spec(dspec, node_spoof=self.node_spoof) return files
class RestrictedBeam(cylinder.CylinderTelescope): beam_height = config.Property(proptype=float, default=30.0) beam_type = config.Property(proptype=str, default="box") def bmask_gaussian(self, feed, freq): pointing = self.zenith bdist = self._angpos - pointing[np.newaxis, :] bdist = np.abs( np.where( (bdist[:, 1] < np.pi)[:, np.newaxis], bdist, bdist - np.array([0, 2 * np.pi])[np.newaxis, :], )) bmask = gaussian_fwhm(bdist[:, 0], np.radians(self.beam_height)) return bmask def bmask_box(self, feed, freq): pointing = self.zenith bdist = self._angpos - pointing[np.newaxis, :] bdist = np.abs( np.where( (bdist[:, 1] < np.pi)[:, np.newaxis], bdist, bdist - np.array([0, 2 * np.pi])[np.newaxis, :], )) bmask = np.abs(bdist[:, 0] / np.radians(self.beam_height)) < 0.5 return bmask
class GradientCylinder(cylinder.UnpolarisedCylinderTelescope): min_spacing = config.Property(proptype=float, default=-1.0) max_spacing = config.Property(proptype=float, default=20.0) def feed_positions_cylinder(self, cylinder_index): if cylinder_index >= self.num_cylinders or cylinder_index < 0: raise Exception("Cylinder index is invalid.") nf = self.num_feeds # Parameters for gradient feedspacing a = self.wavelengths[-1] / 2.0 if self.min_spacing < 0.0 else self.min_spacing # b = 2 * (sp - a) / nf b = 2.0 * (self.max_spacing - a * (nf - 1)) / (nf - 1) ** 2.0 pos = np.empty([nf, 2], dtype=np.float64) i = np.arange(nf) pos[:, 0] = cylinder_index * self.cylinder_spacing pos[:, 1] = a * i + 0.5 * b * i ** 2 return pos
class DummyTask(SingleTask): """Produce an empty data stream for testing. Attributes ---------- total_len : int Length of output data stream. Default: 1. tag : str What to use as a tag for the produced data. """ total_len = config.Property(default=1, proptype=int) tag = config.Property(proptype=str) def process(self): """Produce an empty stream and pass on. Returns ------- cont : subclass of `memh5.BasicCont` Empty data stream. """ if self.total_len == 0: raise pipeline.PipelineStopIteration self.log.debug("Producing test data '{}'...".format(self.tag)) cont = memh5.BasicCont() if "tag" not in cont.attrs: cont.attrs["tag"] = self.tag self.total_len -= 1 return cont
class ApplyScaleFactor(task.SingleTask): """ Apply an overall scale factor to the visibilities. Attributes ---------- scale_factor : float Multiply the visibilities by this number. """ scale_factor = config.Property(proptype=float, default=1e7) update_weight = config.Property(proptype=bool, default=False) def process(self, tstream): """ Parameters ---------- tstream : containers.TimeStream or containers.SiderealStream Returns ------- tstream : containers.TimeStream or containers.SiderealStream """ tstream.vis[:] *= self.scale_factor if self.update_weight: tstream.weight[:] /= self.scale_factor**2 return tstream
class SquareGridLayout(config.Reader): spacing = config.Property(proptype=float, default=6.) grid_size = config.Property(proptype=int, default=3) @util.cache_last def __call__(self): pos_u, pos_v = np.meshgrid( np.linspace(0, self.spacing*(self.grid_size-1), self.grid_size), np.linspace(0, self.spacing*(self.grid_size-1), self.grid_size)) return np.column_stack((pos_u.T.flat, pos_v.T.flat))
class ThresholdVisWeight(task.SingleTask): """Set any weight less than the user specified threshold equal to zero. Threshold is determined as `maximum(absolute_threshold, relative_threshold * mean(weight))` and is evaluated per product/stack entry. Parameters ---------- absolute_threshold : float Any weights with values less than this number will be set to zero. relative_threshold : float Any weights with values less than this number times the average weight will be set to zero. """ absolute_threshold = config.Property(proptype=float, default=1e-7) relative_threshold = config.Property(proptype=float, default=0.0) def process(self, timestream): """Apply threshold to `weight` dataset. Parameters ---------- timestream : `.core.container` with `weight` attribute Returns ------- timestream : same as input timestream The input container with modified weights. """ timestream.redistribute(["prod", "stack"]) weight = timestream.weight[:] # Average over the frequency and time axes to get a per baseline # average mean_weight = weight.mean(axis=2).mean(axis=0) # Figure out which entries to keep threshold = np.maximum(self.absolute_threshold, self.relative_threshold * mean_weight) keep = weight > threshold[np.newaxis, :, np.newaxis] keep_total = timestream.comm.allreduce(np.sum(keep)) keep_frac = keep_total / float(np.prod(weight.global_shape)) self.log.info("%0.5f%% of data is below the weight threshold" % (100.0 * (1.0 - keep_frac))) timestream.weight[:] = np.where(keep, weight, 0.0) return timestream
class AiryBeam(config.Reader): diameter = config.Property(proptype=float, default=6.) sep_limit = config.Property(proptype=float, default=90.) def __call__(self, angpos, zenith, wavelength, feed, pol_index): seps = separations(angpos, zenith) x = np.pi*self.diameter/wavelength*np.sin(seps) out = (2*bessel_j1(x)/x) # (Voltage Beam) out[np.degrees(seps) > self.sep_limit] = 0. return out
class LoadFilesFromParams(pipeline.TaskBase): """Load data from files given in the tasks parameters. Attributes ---------- files : glob pattern, or list Can either be a glob pattern, or lists of actual files. """ files = config.Property(proptype=_list_or_glob) tag_search = config.Property(proptype=str, default=None) def next(self): """Load the given files in turn and pass on. Returns ------- cont : subclass of `memh5.BasicCont` """ from caput import memh5 if len(self.files) == 0: raise pipeline.PipelineStopIteration # Fetch and remove the first item in the list file_ = self.files.pop(0) print "Loading file %s" % file_ cont = memh5.BasicCont.from_file(file_, distributed=True) # Determine an appropriate tag for the container. # First check for tag in attributes. tag = cont.attrs.get('tag', None) # If tag_search keyword has been input, then obtain tag # from keyword search of the file's directories. if self.tag_search is not None: search_res = [tdir for tdir in os.path.normpath(file_).split(os.sep) if tdir.startswith(self.tag_search)] if search_res: tag = '_'.join(search_res) # Otherwise, use the first part of the filename as the tag. if tag is None: tag = os.path.splitext(os.path.basename(file_))[0] cont.attrs['tag'] = tag return cont
class MaskData(task.SingleTask): """Mask out data ahead of map making. Attributes ---------- auto_correlations : bool Exclude auto correlations if set (default=False). m_zero : bool Ignore the m=0 mode (default=False). positive_m : bool Include positive m-modes (default=True). negative_m : bool Include negative m-modes (default=True). """ auto_correlations = config.Property(proptype=bool, default=False) m_zero = config.Property(proptype=bool, default=False) positive_m = config.Property(proptype=bool, default=True) negative_m = config.Property(proptype=bool, default=True) def process(self, mmodes): """Mask out unwanted datain the m-modes. Parameters ---------- mmodes : containers.MModes Returns ------- mmodes : containers.MModes """ mmodes.redistribute("freq") mw = mmodes.weight[:] # Exclude auto correlations if set if not self.auto_correlations: for pi, (fi, fj) in enumerate(mmodes.prodstack): if fi == fj: mw[..., pi] = 0.0 # Apply m based masks if not self.m_zero: mw[0] = 0.0 if not self.positive_m: mw[1:, 0] = 0.0 if not self.negative_m: mw[1:, 1] = 0.0 return mmodes
class EdgeFlagger(task.SingleTask): """Flag the edges of the transit. Parameters ---------- num_begin: int Number of samples to flag at the start of the transit. num_end: int Number of samples to flag at the end of the transit. """ num_begin = config.Property(proptype=int, default=15) num_end = config.Property(proptype=int, default=15) def process(self, track): """Extend the region that has weight set to zero. Parameters ---------- track: draco.core.containers.TrackBeam Holography track to flag. Returns ------- track: draco.core.containers.TrackBeam Holography track with weight set to zero for `num_begin` samples at the start of the transit and `num_end` samples at the end. """ if (self.num_begin == 0) and (self.num_end == 0): return track track.redistribute("freq") weight = track["weight"][:].view(np.ndarray) for ind in np.ndindex(*weight.shape[:-1]): flag = np.flatnonzero(weight[ind] > 0.0) if flag.size > 0: imin, imax = np.percentile(flag, [0, 100]).astype(np.int) imax = imax + 1 if self.num_begin: weight[ind][imin : imin + self.num_begin] = 0.0 if self.num_end: weight[ind][imax - self.num_end : imax] = 0.0 return track
class PolarisedCylinderTelescope(CylinderTelescope, telescope.SimplePolarisedTelescope): """A complete class for an Unpolarised Cylinder telescope. """ # Change the illuminated width in X and Y illumination_x = config.Property(proptype=float, default=1.0) illumination_y = config.Property(proptype=float, default=1.0) ortho_pol = config.Property(proptype=bool, default=True) # @util.cache_last def beamx(self, feed, freq): bpat = visibility.cylinder_beam( self._angpos, self.zenith, self.illumination_x * self.cylinder_width / self.wavelengths[freq], ) bm = np.zeros_like(self._angpos) if self.ortho_pol: bm[:, 1] = bpat else: thatz, phatz = coord.thetaphi_plane_cart(self.zenith) thatp, phatp = coord.thetaphi_plane_cart(self._angpos) bm[:, 0] = np.dot(thatp, phatz) * bpat bm[:, 1] = np.dot(phatp, phatz) * bpat return bm # @util.cache_last def beamy(self, feed, freq): bpat = visibility.cylinder_beam( self._angpos, self.zenith, self.illumination_y * self.cylinder_width / self.wavelengths[freq], ) bm = np.zeros_like(self._angpos) if self.ortho_pol: bm[:, 0] = bpat else: thatz, phatz = coord.thetaphi_plane_cart(self.zenith) thatp, phatp = coord.thetaphi_plane_cart(self._angpos) bm[:, 0] = np.dot(thatp, thatz) * bpat bm[:, 1] = np.dot(phatp, thatz) * bpat return bm
class LoadFilesFromParams(task.SingleTask): """Load data from files given in the tasks parameters. Attributes ---------- files : glob pattern, or list Can either be a glob pattern, or lists of actual files. distributed : bool, optional Whether the file should be loaded distributed across ranks. """ files = config.Property(proptype=_list_or_glob) distributed = config.Property(proptype=bool, default=True) def process(self): """Load the given files in turn and pass on. Returns ------- cont : subclass of `memh5.BasicCont` """ from caput import memh5 # Garbage collect to workaround leaking memory from containers. # TODO: find actual source of leak import gc gc.collect() if len(self.files) == 0: raise pipeline.PipelineStopIteration # Fetch and remove the first item in the list file_ = self.files.pop(0) self.log.info("Loading file %s" % file_) cont = memh5.BasicCont.from_file( file_, distributed=self.distributed, comm=self.comm ) if "tag" not in cont.attrs: # Get the first part of the actual filename and use it as the tag tag = os.path.splitext(os.path.basename(file_))[0] cont.attrs["tag"] = tag return cont
class SaveModuleVersions(task.SingleTask): """Write module versions to a YAML file. The list of modules should be added to the configuration under key 'save_versions'. The version strings are written to a YAML file. Attributes ---------- root : str Root of the file name to output to. """ root = config.Property(proptype=str) done = True def setup(self): """Save module versions.""" fname = "{}_versions.yml".format(self.root) f = open(fname, "w") f.write(yamldump(self.versions)) f.close() self.done = True def process(self): """Do nothing.""" self.done = True return
class LoadSetupFile(io.BaseLoadFiles): """Loads a file from disk into a memh5 container during setup. Attributes ---------- filename : str Path to a saved container. """ filename = config.Property(proptype=str) def setup(self): """Load the file into a container. Returns ------- cont : subclass of `memh5.BasicCont` """ # Call the baseclass setup to resolve any selections super().setup() # Load the requested file cont = self._load_file(self.filename) # Set the done attribute so the pipeline recognizes this task is finished self.done = True return cont def process(self): pass
class ReceiverTemperature(task.SingleTask): """Add a basic receiver temperature term into the data. This class adds in an uncorrelated, frequency and time independent receiver noise temperature to the data. As it is uncorrelated this will only affect the auto-correlations. Note this only adds in the offset to the visibility, to add the corresponding random fluctuations to subsequently use the :class:`SampleNoise` task. Attributes ---------- recv_temp : float The receiver temperature in Kelvin. """ recv_temp = config.Property(proptype=float, default=0.0) def process(self, data): # Iterate over the products to find the auto-correlations and add the noise into them for pi, prod in enumerate(data.index_map['prod']): # Great an auto! if prod[0] == prod[1]: data.vis[:, pi] += self.recv_temp return data
class GaussianBeam(config.Reader): diameter = config.Property(proptype=float, default=6) fwhm_factor = config.Property(proptype=float, default=1.0) sep_limit = config.Property(proptype=float, default=np.inf) def __call__(self, angpos, zenith, wavelength, feed, pol_index): fwhm = self.fwhm_factor*wavelength/self.diameter sigma = gaussian_fwhm_to_sigma*fwhm seps = separations(angpos, zenith) out = np.exp(-seps**2/2/sigma**2)**0.5 # (Voltage Beam) if np.isfinite(self.sep_limit): out[np.degrees(seps) > self.sep_limit] = 0. return out
class LoadProductManager(pipeline.TaskBase): """Loads a driftscan product manager from disk. Attributes ---------- product_directory : str Path to the root of the products. This is the same as the output directory used by ``drift-makeproducts``. """ product_directory = config.Property(proptype=str) def setup(self): """Load the beam transfer matrices. Returns ------- manager : ProductManager Object describing the telescope. """ import os from drift.core import manager if not os.path.exists(self.product_directory): raise RuntimeError("Products do not exist.") # Load ProductManager and Timestream pm = manager.ProductManager.from_config(self.product_directory) return pm
class SaveConfig(task.SingleTask): """Write pipeline config to a text file. Yaml configuration document is written to a text file. Attributes ---------- root : str Root of the file name to output to. """ root = config.Property(proptype=str) done = True def setup(self): """Save module versions.""" fname = "{}_config.yml".format(self.root) f = open(fname, "w") f.write(yamldump(self.pipeline_config)) f.close() self.done = True def process(self): """Do nothing.""" self.done = True return
class HEALPixBeam(config.Reader): """ Much of this is hard-coded and would benefit from a class that handles all of this better... Assumes hdf5 file containts "beam" dataset with dims: pol_index (x/y feed), freq, healpix pixels, v/h Must setup so h5py will read as complex (usually 'r' and 'i' compound dset) Assumes the input beam is pointed at [phi=pi, theta=pi/2] Also an "index_map/freqs" with frequencies in MHz This can fail badly if the beamtransfer class ends up using a different nside... """ filename = config.Property(proptype=str) def __call__(self, angpos, zenith, wavelength, feed, pol_index): freq = (wavelength*units.m).to('MHz', equivalencies=units.spectral()).value with h5py.File(self.filename, 'r') as fil: freq_ind = np.argmin(np.abs(fil['index_map/freqs'][()] - freq)) beam = fil['beam'][pol_index, freq_ind, :][()] rot = [0, zenith[0]-np.pi/2] r = hp.Rotator(deg=False, rot=rot) return np.stack( [r.rotate_map_pixel(beam[..., 0]), r.rotate_map_pixel(beam[..., 1]) ], axis=-1)
class ThresholdVisWeight(task.SingleTask): """Set any weight less than the user specified threshold equal to zero. Threshold is determined as `maximum(absolute_threshold, relative_threshold * mean(weight))`. Parameters ---------- absolute_threshold : float Any weights with values less than this number will be set to zero. relative_threshold : float Any weights with values less than this number times the average weight will be set to zero. """ absolute_threshold = config.Property(proptype=float, default=1e-7) relative_threshold = config.Property(proptype=float, default=0.0) def process(self, timestream): """Apply threshold to `weight` dataset. Parameters ---------- timestream : `.core.container` with `weight` attribute Returns ------- timestream : same as input timestream The input container with modified weights. """ weight = timestream.weight[:] threshold = self.absolute_threshold if self.relative_threshold > 0.0: sum_weight = self.comm.allreduce(np.sum(weight)) mean_weight = sum_weight / float(np.prod(weight.global_shape)) threshold = np.maximum(threshold, self.relative_threshold * mean_weight) keep = weight > threshold self.log.info( "%0.5f%% of data is below the weight threshold of %0.1e." % (100.0 * (1.0 - np.sum(keep) / float(keep.size)), threshold) ) timestream.weight[:] = np.where(keep, weight, 0.0) return timestream
class LayoutFile(config.Reader): filename = config.Property(proptype=str) @util.cache_last def __call__(self): pos_u, pos_v = np.loadtxt(self.filename, unpack=True) return np.column_stack((pos_u, pos_v))
class ChangeSiderealMean(task.SingleTask): """Subtract or add an overall offset (over time) to each visibility. Parameters ---------- add : bool Add the value instead of subtracting. """ add = config.Property(proptype=bool, default=False) def process(self, sstream, mustream): """ Parameters ---------- sstream : andata.CorrData or containers.SiderealStream Timestream or sidereal stream. mustream : andata.CorrData or containers.SiderealStream Timestream or sidereal stream with 1 element in the time axis that contains the value to add or subtract. Returns ------- sstream : same as input Timestream or sidereal stream with value added or subtracted. """ # Check that input visibilities have consistent shapes sshp, mshp = sstream.vis.shape, mustream.vis.shape if np.any(sshp[0:2] != mshp[0:2]): ValueError("Frequency or product axis differ between inputs.") if (len(mshp) != 3) or (mshp[-1] != 1): ValueError( "Mean value has incorrect shape, must be (nfreq, nprod, 1).") # Ensure both inputs are distributed over frequency sstream.redistribute("freq") mustream.redistribute("freq") # Determine indices of autocorrelations prod = mustream.index_map["prod"][mustream.index_map["stack"]["prod"]] not_auto = (prod["input_a"] != prod["input_b"]).astype(np.float32) not_auto = not_auto[np.newaxis, :, np.newaxis] # Add or subtract value to the cross-correlations if self.add: sstream.vis[:] += mustream.vis[:].view(np.ndarray) * not_auto else: sstream.vis[:] -= mustream.vis[:].view(np.ndarray) * not_auto # Set weights to zero if there was no mean sstream.weight[:] *= (mustream.weight[:].view(np.ndarray) > 0.0).astype(sstream.weight.dtype) # Return sidereal stream with modified offset return sstream
class SquareGridLayoutRotated(config.Reader): spacing = config.Property(proptype=float, default=6.) grid_size = config.Property(proptype=int, default=3) rotation_angle = config.Property(proptype=float, default=0) @util.cache_last def __call__(self): pos_u, pos_v = np.meshgrid( np.linspace(0, self.spacing*(self.grid_size-1), self.grid_size), np.linspace(0, self.spacing*(self.grid_size-1), self.grid_size)) temp = np.column_stack((pos_u.T.flat, pos_v.T.flat)) for i in range(len(temp)): temp[i] = rotate_coord(temp[i],self.rotation_angle) return temp
class QueryInputs(task.MPILoggedTask): """From a dataspec describing the data create a list of objects describing the inputs in the files. Attributes ---------- cache : bool Only query for the inputs for the first container received. For all subsequent files just return the initial set of inputs. This can help minimise the number of potentially fragile database operations. """ cache = config.Property(proptype=bool, default=False) _cached_inputs = None def next(self, ts): """Generate an input description from the timestream passed in. Parameters ---------- ts : andata.CorrData Timestream container. Returns ------- inputs : list of :class:`CorrInput` A list of describing the inputs as they are in the file. """ # Fetch from the cache if we can if self.cache and self._cached_inputs: self.log.debug("Using cached inputs.") return self._cached_inputs inputs = None if mpiutil.rank0: # Get the datetime of the middle of the file time = ephemeris.unix_to_datetime(0.5 * (ts.time[0] + ts.time[-1])) inputs = tools.get_correlator_inputs(time) inputs = tools.reorder_correlator_inputs(ts.index_map["input"], inputs) # Broadcast input description to all ranks inputs = mpiutil.world.bcast(inputs, root=0) # Save into the cache for the next iteration if self.cache: self._cached_inputs = inputs # Make sure all nodes have container before return mpiutil.world.Barrier() return inputs
class RandomTask(task.MPILoggedTask): """A base class for MPI tasks that needs to generate random numbers. Attributes ---------- seed : int, optional Set the seed for use in the task. If not set, a random seed is generated and broadcast to all ranks. The seed being used is logged, to repeat a previous run, simply set this as the seed parameter. """ seed = config.Property(proptype=int, default=None) _rng = None @property def rng(self): """A random number generator for this task. .. warning:: Initialising the RNG is a collective operation if the seed is not set, and so all ranks must participate in the first access of this property. Returns ------- rng : np.random.Generator A deterministically seeded random number generator suitable for use in MPI jobs. """ if self._rng is None: # Generate a new base seed for all MPI ranks if self.seed is None: # Use seed sequence to generate a random seed seed = np.random.SeedSequence().entropy seed = self.comm.bcast(seed, root=0) else: seed = self.seed self.log.info("Using random seed: %i", seed) # Construct the new MPI-process and task specific seed. This mixes an # integer checksum of the class name with the MPI-rank to generate a new # hash. # NOTE: the slightly odd (rank + 1) is to ensure that even rank=0 mixes in # the class seed cls_name = "%s.%s" % (self.__module__, self.__class__.__name__) cls_seed = zlib.adler32(cls_name.encode()) new_seed = seed + (self.comm.rank + 1) * cls_seed self._rng = np.random.Generator(_default_bitgen(new_seed)) return self._rng
class SVDSpectrumEstimator(task.SingleTask): """Calculate the SVD spectrum of a set of m-modes. Attributes ---------- niter : int Number of iterations of EM to perform. """ niter = config.Property(proptype=int, default=5) def process(self, mmodes): """Calculate the spectrum. Parameters ---------- mmodes : containers.MModes MModes to find the spectrum of. Returns ------- spectrum : containers.SVDSpectrum """ mmodes.redistribute("m") vis = mmodes.vis[:] weight = mmodes.weight[:] nmode = min(vis.shape[1] * vis.shape[3], vis.shape[2]) spec = containers.SVDSpectrum(singularvalue=nmode, axes_from=mmodes) spec.spectrum[:] = 0.0 for mi, m in vis.enumerate(axis=0): self.log.debug("Calculating SVD spectrum of m=%i", m) vis_m = ( vis[mi].view(np.ndarray).transpose((1, 0, 2)).reshape(vis.shape[2], -1) ) weight_m = ( weight[mi] .view(np.ndarray) .transpose((1, 0, 2)) .reshape(vis.shape[2], -1) ) mask_m = weight_m == 0.0 u, sig, vh = svd_em(vis_m, mask_m, niter=self.niter) spec.spectrum[m] = sig return spec
class SetMPILogging(pipeline.TaskBase): """A task used to configure MPI aware logging. Attributes ---------- level_rank0, level_all : int or str Log level for rank=0, and other ranks respectively. """ level_rank0 = config.Property(proptype=_log_level, default=logging.INFO) level_all = config.Property(proptype=_log_level, default=logging.WARN) def __init__(self): from mpi4py import MPI import math logging.captureWarnings(True) rank_length = int(math.log10(MPI.COMM_WORLD.size)) + 1 mpi_fmt = "[MPI %%(mpi_rank)%id/%%(mpi_size)%id]" % (rank_length, rank_length) filt = MPILogFilter(level_all=self.level_all, level_rank0=self.level_rank0) # This uses the fact that caput.pipeline.Manager has already # attempted to set up the logging. We just override the level, and # insert our custom filter root_logger = logging.getLogger() root_logger.setLevel(logging.DEBUG) ch = root_logger.handlers[0] ch.setLevel(logging.DEBUG) ch.addFilter(filt) formatter = logging.Formatter( "%(elapsedTime)8.1fs " + mpi_fmt + " - %(levelname)-8s %(name)s: %(message)s" ) ch.setFormatter(formatter)
class SmoothVisWeight(task.SingleTask): """Smooth the visibility weights with a median filter. This is done in-place. Attributes ---------- kernel_size : int Size of the kernel for the median filter in time points. Default is 31, corresponding to ~5 minutes window for 10s cadence data. """ # 31 time points correspond to ~ 5min in 10s cadence kernel_size = config.Property(proptype=int, default=31) def process(self, data): """Smooth the weights with a median filter. Parameters ---------- data : :class:`andata.CorrData` or :class:`containers.TimeStream` object Data containing the weights to be smoothed Returns ------- data : Same object as data Data object containing the same data as the input, but with the weights substituted by the smoothed ones. """ # Ensure data is distributed in frequency: data.redistribute("freq") # Full slice reutrns an MPIArray weight = data.weight[:] # Data will be distributed in frequency. # So a frequency loop will not be too large. for lfi, gfi in weight.enumerate(axis=0): # MPIArray takes the local index, returns a local np.ndarray # Find values equal to zero to preserve them in final weights zeromask = weight[lfi] == 0.0 # Median filter. Mode='nearest' to prevent steps close to # the end from being washed weight[lfi] = median_filter(weight[lfi], size=(1, self.kernel_size), mode="nearest") # Ensure zero values are zero weight[lfi][zeromask] = 0.0 return data
class ApplyHolographyGains(task.SingleTask): """Apply gains to a holography transit Attributes ---------- overwrite: bool (default: False) If True, overwrite the input TrackBeam. """ overwrite = config.Property(proptype=bool, default=False) def process(self, track_in, gain): """Apply gain Parameters ---------- track: draco.core.containers.TrackBeam Holography track to apply gains to. Will apply gains to track['beam'], expecting axes to be freq, pol, input, ha gain: np.array Gain to apply. Expected axes are freq, pol, and input Returns ------- track: draco.core.containers.TrackBeam Holography track with gains applied. """ if self.overwrite: track = track_in else: track = TrackBeam( axes_from=track_in, attrs_from=track_in, distributed=track_in.distributed, comm=track_in.comm, ) track["beam"] = track_in["beam"][:] track["weight"] = track_in["weight"][:] track["beam"][:] *= gain.gain[:][:, np.newaxis, :, np.newaxis] track["weight"][:] *= invert_no_zero(np.abs(gain.gain[:]) ** 2)[ :, np.newaxis, :, np.newaxis ] track["beam"][:] = np.where(np.isfinite(track["beam"][:]), track["beam"][:], 0) track["weight"][:] = np.where( np.isfinite(track["weight"][:]), track["weight"][:], 0 ) return track