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