def _resample(self, xgrid, ygrid, wgrid, x):
lza = regrid.lanczos_forward_matrix(xgrid, x, a=self.lanczos_width).T
y = np.matmul(ygrid, lza)
w = invert_no_zero(np.matmul(invert_no_zero(wgrid), lza**2))
return y, w
def process_finish(self):
"""Normalise the stack and return the result.
Includes the sample variance over transits within the stack.
Returns
-------
stack: draco.core.containers.TrackBeam
Stacked transits.
"""
# Divide by norm to get average transit
inv_norm = invert_no_zero(self.norm)
self.stack.beam[:] *= inv_norm
self.stack.weight[:] = invert_no_zero(self.stack.weight[:]) * self.norm**2
self.variance = self.variance * inv_norm - np.abs(self.stack.beam[:]) ** 2
self.pseudo_variance = self.pseudo_variance * inv_norm - self.stack.beam[:] ** 2
# Calculate the covariance between the real and imaginary component
# from the accumulated variance and psuedo-variance
self.stack.sample_variance[0] = 0.5 * (
self.variance + self.pseudo_variance.real
)
self.stack.sample_variance[1] = 0.5 * self.pseudo_variance.imag
self.stack.sample_variance[2] = 0.5 * (
self.variance - self.pseudo_variance.real
)
# Create tag
time_range = np.percentile(self.stack.attrs["transit_time"], [0, 100])
self.stack.attrs["tag"] = "{}_{}_to_{}".format(
self.stack.attrs["source_name"],
ephem.unix_to_datetime(time_range[0]).strftime("%Y%m%dT%H%M%S"),
ephem.unix_to_datetime(time_range[1]).strftime("%Y%m%dT%H%M%S"),
)
return self.stack
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
def process(self, transit):
"""Perform the fit.
Parameters
----------
transit: TrackBeam
Transit to be fit.
Returns
-------
fit: TransitFitParams
Fit parameters.
"""
transit.redistribute("freq")
# Set initial estimate of beam sigma
local_slice = slice(
transit.beam.local_offset[0],
transit.beam.local_offset[0] + transit.beam.local_shape[0],
)
ninput = transit.beam.local_shape[2]
freq = transit.freq[local_slice]
sigma = (0.7 * SPEED_LIGHT / (CHIME_CYL_W * freq)) * (360.0 / np.pi)
sigma = sigma[:, np.newaxis] * np.ones((1, ninput), dtype=sigma.dtype)
# Find index into pol axis that yields copolar products
pol_axis = list(transit.index_map["pol"])
if "co" in pol_axis:
copolar_slice = (slice(None), pol_axis.index("co"))
else:
this_pol = np.array(
[
pol_axis.index("S")
if not ((ii // 256) % 2)
else pol_axis.index("E")
for ii in range(ninput)
]
)
copolar_slice = (slice(None), this_pol, np.arange(ninput))
# Dereference datasets
ha = transit.pix["phi"][:]
vis = transit.beam[:].view(np.ndarray)
vis = vis[copolar_slice]
err = transit.weight[:].view(np.ndarray)
err = np.sqrt(invert_no_zero(err[copolar_slice]))
# Flag data that is outside the fit window set by nsigma config parameter
if self.nsigma is not None:
err *= (
np.abs(ha[np.newaxis, np.newaxis, :])
<= (self.nsigma * sigma[:, :, np.newaxis])
).astype(err.dtype)
# Instantiate the model fitter
model = self.ModelClass(**self.model_kwargs)
# Fit the model
model.fit(ha, vis, err, width=sigma, **self.fit_kwargs)
# Pack into container
fit = TransitFitParams(
param=model.parameter_names,
component=model.component,
axes_from=transit,
attrs_from=transit,
distributed=transit.distributed,
comm=transit.comm,
)
fit.add_dataset("chisq")
fit.add_dataset("ndof")
fit.redistribute("freq")
# Transfer fit information to container attributes
fit.attrs["model_kwargs"] = json.dumps(model.model_kwargs)
fit.attrs["model_class"] = ".".join(
[getattr(self.ModelClass, key) for key in ["__module__", "__name__"]]
)
# Save datasets
fit.parameter[:] = model.param[:]
fit.parameter_cov[:] = model.param_cov[:]
fit.chisq[:] = model.chisq[:]
fit.ndof[:] = model.ndof[:]
return fit
def process(self, beam, data):
"""Stack
Parameters
----------
beam : TrackBeam
The beam that will be stacked.
data : VisContainer
Must contain `prod` index map and `stack` reverse map
that will be used to stack the beam.
Returns
-------
stacked_beam: VisContainer
The input `beam` stacked in the same manner as
"""
# Distribute over frequencies
data.redistribute("freq")
beam.redistribute("freq")
# Grab the stack specifications from the input sidereal stream
prod = data.index_map["prod"]
reverse_stack = data.reverse_map["stack"][:]
input_flags = data.input_flags[:]
if not np.any(input_flags):
input_flags = np.ones_like(input_flags)
# Create output container
if isinstance(data, SiderealStream):
OutputContainer = SiderealStream
output_kwargs = {"ra": data.ra[:]}
else:
OutputContainer = TimeStream
output_kwargs = {"time": data.time[:]}
stacked_beam = OutputContainer(
axes_from=data,
attrs_from=beam,
distributed=True,
comm=data.comm,
**output_kwargs
)
stacked_beam.vis[:] = 0.0
stacked_beam.weight[:] = 0.0
stacked_beam.attrs["tag"] = "_".join([beam.attrs["tag"], data.attrs["tag"]])
# Dereference datasets
bv = beam.beam[:].view(np.ndarray)
bw = beam.weight[:].view(np.ndarray)
ov = stacked_beam.vis[:]
ow = stacked_beam.weight[:]
pol_filter = {
"X": "X",
"Y": "Y",
"E": "X",
"S": "Y",
"co": "co",
"cross": "cross",
}
pol = [pol_filter.get(pp, None) for pp in self.telescope.polarisation]
beam_pol = [pol_filter.get(pp, None) for pp in beam.index_map["pol"][:]]
# Compute the fractional variance of the beam measurement
frac_var = invert_no_zero(bw * np.abs(bv) ** 2)
# Create counter to increment during the stacking.
# This will be used to normalize at the end.
counter = np.zeros_like(ow)
# Construct stack
for pp, (ss, conj) in enumerate(reverse_stack):
aa, bb = prod[pp]
if conj:
aa, bb = bb, aa
try:
aa_pol, bb_pol = self._resolve_pol(pol[aa], pol[bb], beam_pol)
except ValueError:
continue
cross = bv[:, aa_pol, aa, :] * bv[:, bb_pol, bb, :].conj()
weight = (
input_flags[np.newaxis, aa, :]
* input_flags[np.newaxis, bb, :]
* invert_no_zero(
np.abs(cross) ** 2
* (frac_var[:, aa_pol, aa, :] + frac_var[:, bb_pol, bb, :])
)
)
if self.weight == "inverse_variance":
wss = weight
else:
wss = (weight > 0.0).astype(np.float32)
# Accumulate variances in quadrature. Save in the weight dataset.
ov[:, ss, :] += wss * cross
ow[:, ss, :] += wss**2 * invert_no_zero(weight)
# Increment counter
counter[:, ss, :] += wss
# Divide through by counter to get properly weighted visibility average
ov[:] *= invert_no_zero(counter)
ow[:] = counter**2 * invert_no_zero(ow[:])
return stacked_beam
def process(self, transit):
"""Add a transit to the stack.
Parameters
----------
transit: draco.core.containers.TrackBeam
A holography transit.
"""
self.log.info("Weight is %s" % self.weight)
if self.stack is None:
self.log.info("Initializing transit stack.")
self.stack = TrackBeam(
axes_from=transit, distributed=transit.distributed, comm=transit.comm
)
self.stack.add_dataset("sample_variance")
self.stack.add_dataset("nsample")
self.stack.redistribute("freq")
self.log.info("Adding %s to stack." % transit.attrs["tag"])
# Copy over relevant attributes
self.stack.attrs["filename"] = [transit.attrs["tag"]]
self.stack.attrs["observation_id"] = [transit.attrs["observation_id"]]
self.stack.attrs["transit_time"] = [transit.attrs["transit_time"]]
self.stack.attrs["archivefiles"] = list(transit.attrs["archivefiles"])
self.stack.attrs["dec"] = transit.attrs["dec"]
self.stack.attrs["source_name"] = transit.attrs["source_name"]
self.stack.attrs["icrs_ra"] = transit.attrs["icrs_ra"]
self.stack.attrs["cirs_ra"] = transit.attrs["cirs_ra"]
# Copy data for first transit
flag = (transit.weight[:] > 0.0).astype(np.int)
if self.weight == "inverse_variance":
coeff = transit.weight[:]
else:
coeff = flag.astype(np.float32)
self.stack.beam[:] = coeff * transit.beam[:]
self.stack.weight[:] = (coeff**2) * invert_no_zero(transit.weight[:])
self.stack.nsample[:] = flag.astype(np.int)
self.variance = coeff * np.abs(transit.beam[:]) ** 2
self.pseudo_variance = coeff * transit.beam[:] ** 2
self.norm = coeff
else:
if list(transit.beam.shape) != list(self.stack.beam.shape):
self.log.error(
"Transit has different shape than stack: {}, {}".format(
transit.beam.shape, self.stack.beam.shape
)
+ " Skipping."
)
return None
self.log.info("Adding %s to stack." % transit.attrs["tag"])
self.stack.attrs["filename"].append(transit.attrs["tag"])
self.stack.attrs["observation_id"].append(transit.attrs["observation_id"])
self.stack.attrs["transit_time"].append(transit.attrs["transit_time"])
self.stack.attrs["archivefiles"] += list(transit.attrs["archivefiles"])
# Accumulate transit data
flag = (transit.weight[:] > 0.0).astype(np.int)
if self.weight == "inverse_variance":
coeff = transit.weight[:]
else:
coeff = flag.astype(np.float32)
self.stack.beam[:] += coeff * transit.beam[:]
self.stack.weight[:] += (coeff**2) * invert_no_zero(transit.weight[:])
self.stack.nsample[:] += flag
self.variance += coeff * np.abs(transit.beam[:]) ** 2
self.pseudo_variance += coeff * transit.beam[:] ** 2
self.norm += coeff
return None
def process(self, sstream):
"""Computes the ringmap.
Parameters
----------
sstream : containers.SiderealStream
The input sidereal stream.
Returns
-------
rm : containers.RingMap
"""
# Redistribute over frequency
sstream.redistribute("freq")
nfreq = sstream.vis.local_shape[0]
# Extract the right ascension (or calculate from timestamp)
ra = sstream.ra if "ra" in sstream.index_map else ephemeris.lsa(sstream.time)
nra = ra.size
# Construct mapping from vis array to unpacked 2D grid
nprod = sstream.prodstack.shape[0]
pind = np.zeros(nprod, dtype=np.int)
xind = np.zeros(nprod, dtype=np.int)
ysep = np.zeros(nprod, dtype=np.float)
for pp, (ii, jj) in enumerate(sstream.prodstack):
if self.telescope.feedconj[ii, jj]:
ii, jj = jj, ii
fi = self.telescope.feeds[ii]
fj = self.telescope.feeds[jj]
pind[pp] = 2 * int(fi.pol == "S") + int(fj.pol == "S")
xind[pp] = np.abs(fi.cyl - fj.cyl)
# TODO: don't use this internal property. This can probably wait
# until the RingMapMaker gets completely moved into draco
ysep[pp] = fi._pos[1] - fj._pos[1]
abs_ysep = np.abs(ysep)
min_ysep, max_ysep = np.percentile(abs_ysep[abs_ysep > 0.0], [0, 100])
yind = np.round(ysep / min_ysep).astype(np.int)
grid_index = list(zip(pind, xind, yind))
# Define several variables describing the baseline configuration.
nfeed = int(np.round(max_ysep / min_ysep)) + 1
nvis_1d = 2 * nfeed - 1
ncyl = np.max(xind) + 1
nbeam = 1 if self.single_beam else int(2 * ncyl - 1)
# Define polarisation axis
pol = np.array(["XX", "reXY", "imXY", "YY"])
npol = len(pol)
# Create empty array for output
vis = np.zeros((nfreq, npol, nra, ncyl, nvis_1d), dtype=np.complex128)
invvar = np.zeros((nfreq, npol, nra, ncyl, nvis_1d), dtype=np.float64)
weight = np.zeros((nfreq, npol, nra, ncyl, nvis_1d), dtype=np.float64)
# If natural or uniform weighting was chosen, then calculate the
# redundancy of the collated visibilities.
if self.weight != "inverse_variance":
redundancy = tools.calculate_redundancy(
sstream.input_flags[:],
sstream.index_map["prod"][:],
sstream.reverse_map["stack"]["stack"][:],
sstream.vis.shape[1],
)
if self.weight == "uniform":
redundancy = (redundancy > 0).astype(np.float32)
# De-reference distributed arrays outside loop to save repeated MPI calls
ssv = sstream.vis[:]
ssw = sstream.weight[:]
# Unpack visibilities into new array
for vis_ind, (p_ind, x_ind, y_ind) in enumerate(grid_index):
# Handle different options for weighting baselines
if self.weight == "inverse_variance":
w = ssw[:, vis_ind]
else:
w = (ssw[:, vis_ind] > 0.0).astype(np.float32)
w *= redundancy[np.newaxis, vis_ind]
if x_ind != 0 or not self.exclude_intracyl:
vis[:, p_ind, :, x_ind, y_ind] = ssv[:, vis_ind]
invvar[:, p_ind, :, x_ind, y_ind] = ssw[:, vis_ind]
weight[:, p_ind, :, x_ind, y_ind] = w
# Remove auto-correlations
if not self.include_auto:
weight[..., 0, 0] = 0.0
# Autos get double-counted at the end
weight[..., 0, 0] *= 0.5
# Normalize the weighting function
# Multiply by 2 here to count negative baselines
norm = 2 * np.sum(weight, axis=(-2, -1))
weight *= tools.invert_no_zero(norm)[..., np.newaxis, np.newaxis]
# Construct phase array
el = self.span * np.linspace(-1.0, 1.0, self.npix)
vis_pos_1d = np.fft.fftfreq(nvis_1d, d=(1.0 / (nvis_1d * min_ysep)))
# Create empty ring map
rm = containers.RingMap(
beam=nbeam, el=el, pol=pol, ra=ra, axes_from=sstream, attrs_from=sstream
)
# Add datasets
rm.add_dataset("rms")
rm.add_dataset("dirty_beam")
# Make sure ring map is distributed over frequency
rm.redistribute("freq")
# Estimate rms noise in the ring map by propagating estimates
# of the variance in the visibilities
rm.rms[:] = np.sqrt(
2 * np.sum(tools.invert_no_zero(invvar) * weight**2.0, axis=(-2, -1))
).transpose(1, 0, 2)
# Dereference datasets
rmm = rm.map[:]
rmb = rm.dirty_beam[:]
# Pre-allocate arrays that will be reused inside loop
save_nb = 1 if self.single_beam else nbeam
pa = np.zeros(vis_pos_1d.shape + el.shape, dtype=vis.dtype)
bfm_y = np.zeros((npol, nra, ncyl, self.npix), dtype=vis.dtype)
sb_y = np.zeros((npol, nra, ncyl, self.npix), dtype=vis.dtype)
bfm = np.zeros((npol, nra, save_nb, self.npix), dtype=vis.dtype)
sb = np.zeros((npol, nra, save_nb, self.npix), dtype=vis.dtype)
# Loop over local frequencies and fill ring map
for lfi, fi in sstream.vis[:].enumerate(0):
# Get the current frequency and wavelength
fr = sstream.freq[fi]
wv = scipy.constants.c * 1e-6 / fr
# Inverse discrete Fourier transform in y-direction
pa[:] = np.exp(
(-2.0j * np.pi / wv * vis_pos_1d[:, np.newaxis]) * el[np.newaxis, :]
)
# Perform inverse discrete fourier transform in y-direction
# and inverse fast fourier transform in x-direction
bfm_y[:] = np.matmul(weight[lfi] * vis[lfi], pa)
sb_y[:] = np.matmul(weight[lfi], pa)
if self.single_beam:
# Only need the 0th term if the irfft, equivalent to adding in EW
# direction
bfm[:] = np.sum(bfm_y, axis=2)[:, :, np.newaxis, ...]
sb[:] = np.sum(sb_y, axis=2)[:, :, np.newaxis, ...]
else:
bfm[:] = np.fft.ifft(bfm_y, nbeam, axis=2) * nbeam
sb[:] = np.fft.ifft(sb_y, nbeam, axis=2) * nbeam
# Save to container (shifting to the final axis ordering)
# for co-pol we take twice the real part
# to complete sum over negative baselines
copol_ind = [0, 3]
rmm[:, copol_ind, lfi] = 2 * bfm[copol_ind].real.transpose(2, 0, 1, 3)
rmb[:, copol_ind, lfi] = 2 * sb[copol_ind].real.transpose(2, 0, 1, 3)
# for cross-pol we save real and imaginary parts of complex map formed
# by combining positive and negative baselines (which are in the other index)
xpol_ind = [1, 2]
rmm[:, xpol_ind, lfi] = (
(bfm[xpol_ind[0]] + bfm[xpol_ind[1]].conj())
.view("(2,)float")
.transpose(1, 3, 0, 2)
)
rmb[:, xpol_ind, lfi] = (
(sb[xpol_ind[0]] + sb[xpol_ind[1]].conj())
.view("(2,)float")
.transpose(1, 3, 0, 2)
)
return rm
def process(self, tstream, inputmap):
"""Apply the decorrelation correction for a given source.
Parameters
----------
tstream : andata.CorrData
timestream data
inputmap : list of :class:`CorrInput`
A list of describing the inputs as they are in the file, output from
`tools.get_correlator_inputs()`.
Returns
-------
tstream : andata.CorrData
Returns the corrected timestream.
"""
tstream.redistribute("freq")
prod_map = tstream.prodstack
src = ephemeris.source_dictionary[self.source]
# Rotate the telescope
tools.change_chime_location(rotation=self.telescope_rotation)
# correct visibilities
corr_vis = tools.decorrelation(
tstream.vis[:],
times=tstream.time,
feeds=inputmap,
src=src,
prod_map=prod_map,
wterm=self.wterm,
inplace=self.overwrite,
)
# weights are inverse variance
weight = np.sqrt(invert_no_zero(tstream.weight[:]))
weight = tools.decorrelation(
weight,
times=tstream.time,
feeds=inputmap,
src=src,
prod_map=prod_map,
wterm=self.wterm,
inplace=True,
)
weight = invert_no_zero(weight) ** 2
# Return telescope to default rotation
tools.change_chime_location(default=True)
if self.overwrite:
# visibilities were corrected in place
tstream.weight[:] = weight
return tstream
else:
tstream_corr = containers.empty_like(tstream)
tstream_corr.vis[:] = corr_vis
tstream_corr.weight[:] = weight
return tstream_corr