def gen_lm(sel_PHASE_DIR, sel_LSM, n_dir):
lm = np.zeros((n_dir, 2), dtype=np.float64)
# Cycle coordinates creating a source with flux
for d, source in enumerate(sel_LSM.sources):
# Extract position
radec_s = np.array([[source.pos.ra, source.pos.dec]])
lm[d] = radec_to_lm(radec_s, sel_PHASE_DIR)
return lm
def new(ms, sky_model, gains, **kwargs):
"""Generate model visibilties per source (as direction axis)
for stokes I and Q and generate relevant visibilities."""
# Options to attributed dictionary
if kwargs["yaml"] is not None:
options = ocf.load(kwargs["yaml"])
else:
options = ocf.create(kwargs)
# Set to struct
ocf.set_struct(options, True)
# Change path to sky model if chosen
try:
sky_model = sky_models[sky_model.lower()]
except:
# Own sky model reference
pass
# Set thread count to cpu count
if options.ncpu:
from multiprocessing.pool import ThreadPool
import dask
dask.config.set(pool=ThreadPool(options.ncpu))
else:
import multiprocessing
options.ncpu = multiprocessing.cpu_count()
# Load gains to corrupt with
with open(gains, "rb") as file:
jones = np.load(file)
# Load dimensions
n_time, n_ant, n_chan, n_dir, n_corr = jones.shape
n_row = n_time * (n_ant * (n_ant - 1) // 2)
# Load ms
MS = xds_from_ms(ms)[0]
# Get time-bin indices and counts
row_chunks, tbin_indices, tbin_counts = chunkify_rows(
MS.TIME, options.utime)
# Close and reopen with chunked rows
MS.close()
MS = xds_from_ms(ms, chunks={"row": row_chunks})[0]
# Get antenna arrays (dask ignored for now)
ant1 = MS.ANTENNA1.data
ant2 = MS.ANTENNA2.data
# Adjust UVW based on phase-convention
if options.phase_convention.upper() == 'CASA':
uvw = -MS.UVW.data.astype(np.float64)
elif options.phase_convention.upper() == 'CODEX':
uvw = MS.UVW.data.astype(np.float64)
else:
raise ValueError("Unknown sign convention for phase.")
# MS dimensions
dims = ocf.create(dict(MS.sizes))
# Close MS
MS.close()
# Build source model from lsm
lsm = Tigger.load(sky_model)
# Check if dimensions match jones
assert n_time * (n_ant * (n_ant - 1) // 2) == dims.row
assert n_time == len(tbin_indices)
assert n_ant == np.max((np.max(ant1), np.max(ant2))) + 1
assert n_chan == dims.chan
assert n_corr == dims.corr
# If gains are DIE
if options.die:
assert n_dir == 1
n_dir = len(lsm.sources)
else:
assert n_dir == len(lsm.sources)
# Get phase direction
radec0_table = xds_from_table(ms + '::FIELD')[0]
radec0 = radec0_table.PHASE_DIR.data.squeeze().compute()
radec0_table.close()
# Get frequency column
freq_table = xds_from_table(ms + '::SPECTRAL_WINDOW')[0]
freq = freq_table.CHAN_FREQ.data.astype(np.float64)[0]
freq_table.close()
# Get feed orientation
feed_table = xds_from_table(ms + '::FEED')[0]
feeds = feed_table.POLARIZATION_TYPE.data[0].compute()
# Create initial model array
model = np.zeros((n_dir, n_chan, n_corr), dtype=np.float64)
# Create initial coordinate array and source names
lm = np.zeros((n_dir, 2), dtype=np.float64)
source_names = []
# Cycle coordinates creating a source with flux
print("==> Building model visibilities")
for d, source in enumerate(lsm.sources):
# Extract name
source_names.append(source.name)
# Extract position
radec_s = np.array([[source.pos.ra, source.pos.dec]])
lm[d] = radec_to_lm(radec_s, radec0)
# Get flux - Stokes I
if source.flux.I:
I0 = source.flux.I
# Get spectrum (only spi currently supported)
tmp_spec = source.spectrum
spi = [tmp_spec.spi if tmp_spec is not None else 0.0]
ref_freq = [tmp_spec.freq0 if tmp_spec is not None else 1.0]
# Generate model flux
model[d, :, 0] = I0 * (freq / ref_freq)**spi
# Get flux - Stokes Q
if source.flux.Q:
Q0 = source.flux.Q
# Get spectrum
tmp_spec = source.spectrum
spi = [tmp_spec.spi if tmp_spec is not None else 0.0]
ref_freq = [tmp_spec.freq0 if tmp_spec is not None else 1.0]
# Generate model flux
model[d, :, 1] = Q0 * (freq / ref_freq)**spi
# Get flux - Stokes U
if source.flux.U:
U0 = source.flux.U
# Get spectrum
tmp_spec = source.spectrum
spi = [tmp_spec.spi if tmp_spec is not None else 0.0]
ref_freq = [tmp_spec.freq0 if tmp_spec is not None else 1.0]
# Generate model flux
model[d, :, 2] = U0 * (freq / ref_freq)**spi
# Get flux - Stokes V
if source.flux.V:
V0 = source.flux.V
# Get spectrum
tmp_spec = source.spectrum
spi = [tmp_spec.spi if tmp_spec is not None else 0.0]
ref_freq = [tmp_spec.freq0 if tmp_spec is not None else 1.0]
# Generate model flux
model[d, :, 3] = V0 * (freq / ref_freq)**spi
# Close sky-model
del lsm
# Build dask graph
tbin_indices = da.from_array(tbin_indices, chunks=(options.utime))
tbin_counts = da.from_array(tbin_counts, chunks=(options.utime))
lm = da.from_array(lm, chunks=lm.shape)
model = da.from_array(model, chunks=model.shape)
jones = da.from_array(jones, chunks=(options.utime, ) + jones.shape[1::])
# Apply image to visibility for each source
sources = []
for s in range(n_dir):
source_vis = im_to_vis(model[s].reshape((1, n_chan, n_corr)),
uvw,
lm[s].reshape((1, 2)),
freq,
dtype=np.complex64,
convention='fourier')
sources.append(source_vis)
model_vis = da.stack(sources, axis=2)
# Sum over direction?
if options.die:
model_vis = da.sum(model_vis, axis=2, keepdims=True)
n_dir = 1
source_names = [options.mname]
# Select schema based on feed orientation
if (feeds == ["X", "Y"]).all():
out_schema = [["XX", "XY"], ["YX", "YY"]]
elif (feeds == ["R", "L"]).all():
out_schema = [['RR', 'RL'], ['LR', 'LL']]
else:
raise ValueError("Unknown feed orientation implementation.")
# Convert Stokes to Correlations
in_schema = ['I', 'Q', 'U', 'V']
model_vis = convert(model_vis, in_schema, out_schema).reshape(
(n_row, n_chan, n_dir, n_corr))
# Apply gains to model_vis
print("==> Corrupting visibilities")
data = corrupt_vis(tbin_indices, tbin_counts, ant1, ant2, jones, model_vis)
# Reopen MS
MS = xds_from_ms(ms, chunks={"row": row_chunks})[0]
# Assign model visibilities
out_names = []
for d in range(n_dir):
MS = MS.assign(
**{
source_names[d]: (("row", "chan", "corr"),
model_vis[:, :, d].astype(np.complex64))
})
out_names += [source_names[d]]
# Assign noise free visibilities to 'CLEAN_DATA'
MS = MS.assign(
**{
'CLEAN_' + options.dname: (("row", "chan", "corr"),
data.astype(np.complex64))
})
out_names += ['CLEAN_' + options.dname]
# Get noise realisation
if options.std > 0.0:
# Noise matrix
print(f"==> Applying noise (std={options.std}) to visibilities")
noise = []
for i in range(2):
real = da.random.normal(loc=0.0,
scale=options.std,
size=(n_row, n_chan),
chunks=(row_chunks, n_chan))
imag = 1.0j * (da.random.normal(loc=0.0,
scale=options.std,
size=(n_row, n_chan),
chunks=(row_chunks, n_chan)))
noise.append(real + imag)
# Zero matrix for off-diagonals
zero = da.zeros((n_row, n_chan), chunks=(row_chunks, n_chan))
noise.insert(1, zero)
noise.insert(2, zero)
# NP to Dask
noise = da.stack(noise, axis=2).rechunk((row_chunks, n_chan, n_corr))
# Assign noise to 'NOISE'
MS = MS.assign(
**{'NOISE': (("row", "chan", "corr"), noise.astype(np.complex64))})
out_names += ['NOISE']
# Add noise to data and assign to 'DATA'
noisy_data = data + noise
MS = MS.assign(
**{
options.dname: (("row", "chan", "corr"),
noisy_data.astype(np.complex64))
})
out_names += [options.dname]
# Create a write to the table
write = xds_to_table(MS, ms, out_names)
# Submit all graph computations in parallel
print(f"==> Executing `dask-ms` write to `{ms}` for the following columns: "\
+ f"{', '.join(out_names)}")
with ProgressBar():
write.compute()
print(f"==> Completed.")
def simulate(args):
# get full time column and compute row chunks
ms = table(args.ms)
time = ms.getcol('TIME')
row_chunks, tbin_idx, tbin_counts = chunkify_rows(time,
args.utimes_per_chunk)
# convert to dask arrays
tbin_idx = da.from_array(tbin_idx, chunks=(args.utimes_per_chunk))
tbin_counts = da.from_array(tbin_counts, chunks=(args.utimes_per_chunk))
n_time = tbin_idx.size
ant1 = ms.getcol('ANTENNA1')
ant2 = ms.getcol('ANTENNA2')
n_ant = np.maximum(ant1.max(), ant2.max()) + 1
flag = ms.getcol("FLAG")
n_row, n_freq, n_corr = flag.shape
if n_corr == 4:
model_corr = (2, 2)
jones_corr = (2, )
elif n_corr == 2:
model_corr = (2, )
jones_corr = (2, )
elif n_corr == 1:
model_corr = (1, )
jones_corr = (1, )
else:
raise RuntimeError("Invalid number of correlations")
ms.close()
# get phase dir
radec0 = table(args.ms + '::FIELD').getcol('PHASE_DIR').squeeze()
# get freqs
freq = table(args.ms + '::SPECTRAL_WINDOW').getcol('CHAN_FREQ')[0].astype(
np.float64)
assert freq.size == n_freq
# get source coordinates from lsm
lsm = Tigger.load(args.sky_model)
radec = []
stokes = []
spi = []
ref_freqs = []
for source in lsm.sources:
radec.append([source.pos.ra, source.pos.dec])
stokes.append([source.flux.I])
tmp_spec = source.spectrum
spi.append([tmp_spec.spi if tmp_spec is not None else 0.0])
ref_freqs.append([tmp_spec.freq0 if tmp_spec is not None else 1.0])
n_dir = len(stokes)
radec = np.asarray(radec)
lm = radec_to_lm(radec, radec0)
# load in the model file
model = np.zeros((n_freq, n_dir) + model_corr)
stokes = np.asarray(stokes)
ref_freqs = np.asarray(ref_freqs)
spi = np.asarray(spi)
for d in range(n_dir):
Stokes_I = stokes[d] * (freq / ref_freqs[d])**spi[d]
if n_corr == 4:
model[:, d, 0, 0] = Stokes_I
model[:, d, 1, 1] = Stokes_I
elif n_corr == 2:
model[:, d, 0] = Stokes_I
model[:, d, 1] = Stokes_I
else:
model[:, d, 0] = Stokes_I
# append antenna columns
cols = []
cols.append('ANTENNA1')
cols.append('ANTENNA2')
cols.append('UVW')
# load in gains
jones, alphas = make_screen(lm, freq, n_time, n_ant, jones_corr[0])
jones = jones.astype(np.complex128)
jones_shape = jones.shape
jones_da = da.from_array(jones,
chunks=(args.utimes_per_chunk, ) +
jones_shape[1::])
freqs = da.from_array(freq, chunks=(n_freq))
lm = da.from_array(np.tile(lm[None], (n_time, 1, 1)),
chunks=(args.utimes_per_chunk, n_dir, 2))
# change model to dask array
tmp_shape = (n_time, )
for i in range(len(model.shape)):
tmp_shape += (1, )
model = da.from_array(np.tile(model[None], tmp_shape),
chunks=(args.utimes_per_chunk, ) + model.shape)
# load data in in chunks and apply gains to each chunk
xds = xds_from_ms(args.ms, columns=cols, chunks={"row": row_chunks})[0]
ant1 = xds.ANTENNA1.data
ant2 = xds.ANTENNA2.data
uvw = xds.UVW.data
# apply gains
data = compute_and_corrupt_vis(tbin_idx, tbin_counts, ant1, ant2, jones_da,
model, uvw, freqs, lm)
# Assign visibilities to args.out_col and write to ms
xds = xds.assign(
**{
args.out_col: (("row", "chan", "corr"),
data.reshape(n_row, n_freq, n_corr))
})
# Create a write to the table
write = xds_to_table(xds, args.ms, [args.out_col])
# Submit all graph computations in parallel
with ProgressBar():
write.compute()
return jones, alphas
def calibrate(args, jones, alphas):
# simple calibration to test if simulation went as expected.
# Note do not run on large data set
# load data
ms = table(args.ms)
time = ms.getcol('TIME')
_, tbin_idx, tbin_counts = chunkify_rows(time, args.utimes_per_chunk)
n_time = tbin_idx.size
ant1 = ms.getcol('ANTENNA1')
ant2 = ms.getcol('ANTENNA2')
n_ant = np.maximum(ant1.max(), ant2.max()) + 1
uvw = ms.getcol('UVW').astype(np.float64)
data = ms.getcol(args.out_col) # this is where we put the data
# we know it is pure Stokes I so we can solve using diagonals only
data = data[:, :, (0, 3)].astype(np.complex128)
n_row, n_freq, n_corr = data.shape
flag = ms.getcol('FLAG')
flag = flag[:, :, (0, 3)]
# get phase dir
radec0 = table(args.ms + '::FIELD').getcol('PHASE_DIR').squeeze().astype(
np.float64)
# get freqs
freq = table(args.ms + '::SPECTRAL_WINDOW').getcol('CHAN_FREQ')[0].astype(
np.float64)
assert freq.size == n_freq
# now get the model
# get source coordinates from lsm
lsm = Tigger.load(args.sky_model)
radec = []
stokes = []
spi = []
ref_freqs = []
for source in lsm.sources:
radec.append([source.pos.ra, source.pos.dec])
stokes.append([source.flux.I])
tmp_spec = source.spectrum
spi.append([tmp_spec.spi if tmp_spec is not None else 0.0])
ref_freqs.append([tmp_spec.freq0 if tmp_spec is not None else 1.0])
n_dir = len(stokes)
radec = np.asarray(radec)
lm = radec_to_lm(radec, radec0)
# get model visibilities
model = np.zeros((n_row, n_freq, n_dir, 2), dtype=np.complex)
stokes = np.asarray(stokes)
ref_freqs = np.asarray(ref_freqs)
spi = np.asarray(spi)
for d in range(n_dir):
Stokes_I = stokes[d] * (freq / ref_freqs[d])**spi[d]
model[:, :, d, 0:1] = im_to_vis(Stokes_I[None, :, None], uvw,
lm[d:d + 1], freq)
model[:, :, d, 1] = model[:, :, d, 0]
# set weights to unity
weight = np.ones_like(data, dtype=np.float64)
# initialise gains
jones0 = np.ones((n_time, n_ant, n_freq, n_dir, n_corr),
dtype=np.complex128)
# calibrate
ti = timeit()
jones_hat, jhj, jhr, k = gauss_newton(tbin_idx,
tbin_counts,
ant1,
ant2,
jones0,
data,
flag,
model,
weight,
tol=1e-5,
maxiter=100)
print("%i iterations took %fs" % (k, timeit() - ti))
# verify result
for p in range(2):
for q in range(p):
diff_true = np.angle(jones[:, p] * jones[:, q].conj())
diff_hat = np.angle(jones_hat[:, p] * jones_hat[:, q].conj())
try:
assert_array_almost_equal(diff_true, diff_hat, decimal=2)
except Exception as e:
print(e)
def both(args):
"""Generate model data, corrupted visibilities and
gains (phase-only or normal)"""
# Set thread count to cpu count
if args.ncpu:
from multiprocessing.pool import ThreadPool
import dask
dask.config.set(pool=ThreadPool(args.ncpu))
else:
import multiprocessing
args.ncpu = multiprocessing.cpu_count()
# Get full time column and compute row chunks
ms = xds_from_table(args.ms)[0]
row_chunks, tbin_idx, tbin_counts = chunkify_rows(
ms.TIME, args.utimes_per_chunk)
# Convert time rows to dask arrays
tbin_idx = da.from_array(tbin_idx,
chunks=(args.utimes_per_chunk))
tbin_counts = da.from_array(tbin_counts,
chunks=(args.utimes_per_chunk))
# Time axis
n_time = tbin_idx.size
# Get antenna columns
ant1 = ms.ANTENNA1.data
ant2 = ms.ANTENNA2.data
# No. of antennas axis
n_ant = (np.maximum(ant1.max(), ant2.max()) + 1).compute()
# Get flag column
flag = ms.FLAG.data
# Get convention
if args.phase_convention == 'CASA':
uvw = -(ms.UVW.data.astype(np.float64))
elif args.phase_convention == 'CODEX':
uvw = ms.UVW.data.astype(np.float64)
else:
raise ValueError("Unknown sign convention for phase")
# Get rest of dimensions
n_row, n_freq, n_corr = flag.shape
# Raise error if correlation axis too small
if n_corr != 4:
raise NotImplementedError("Only 4 correlations "\
+ "currently supported")
# Get phase direction
radec0_table = xds_from_table(args.ms+'::FIELD')[0]
radec0 = radec0_table.PHASE_DIR.data.squeeze().compute()
# Get frequency column
freq_table = xds_from_table(args.ms+'::SPECTRAL_WINDOW')[0]
freq = freq_table.CHAN_FREQ.data.astype(np.float64)[0]
# Check dimension
assert freq.size == n_freq
# Check for sky-model
if args.sky_model == 'MODEL-1.txt':
args.sky_model = MODEL_1
elif args.sky_model == 'MODEL-4.txt':
args.sky_model = MODEL_4
elif args.sky_model == 'MODEL-50.txt':
args.sky_model = MODEL_50
else:
raise NotImplemented(f"Sky-model {args.sky_model} not in "\
+ "kalcal/datasets/sky_model/")
# Build source model from lsm
lsm = Tigger.load(args.sky_model)
# Direction axis
n_dir = len(lsm.sources)
# Create initial model array
model = np.zeros((n_dir, n_freq, n_corr), dtype=np.float64)
# Create initial coordinate array and source names
lm = np.zeros((n_dir, 2), dtype=np.float64)
source_names = []
# Cycle coordinates creating a source with flux
for d, source in enumerate(lsm.sources):
# Extract name
source_names.append(source.name)
# Extract position
radec_s = np.array([[source.pos.ra, source.pos.dec]])
lm[d] = radec_to_lm(radec_s, radec0)
# Get flux - Stokes I
if source.flux.I:
I0 = source.flux.I
# Get spectrum (only spi currently supported)
tmp_spec = source.spectrum
spi = [tmp_spec.spi if tmp_spec is not None else 0.0]
ref_freq = [tmp_spec.freq0 if tmp_spec is not None else 1.0]
# Generate model flux
model[d, :, 0] = I0 * (freq/ref_freq)**spi
# Get flux - Stokes Q
if source.flux.Q:
Q0 = source.flux.Q
# Get spectrum
tmp_spec = source.spectrum
spi = [tmp_spec.spi if tmp_spec is not None else 0.0]
ref_freq = [tmp_spec.freq0 if tmp_spec is not None else 1.0]
# Generate model flux
model[d, :, 1] = Q0 * (freq/ref_freq)**spi
# Get flux - Stokes U
if source.flux.U:
U0 = source.flux.U
# Get spectrum
tmp_spec = source.spectrum
spi = [tmp_spec.spi if tmp_spec is not None else 0.0]
ref_freq = [tmp_spec.freq0 if tmp_spec is not None else 1.0]
# Generate model flux
model[d, :, 2] = U0 * (freq/ref_freq)**spi
# Get flux - Stokes V
if source.flux.V:
V0 = source.flux.V
# Get spectrum
tmp_spec = source.spectrum
spi = [tmp_spec.spi if tmp_spec is not None else 0.0]
ref_freq = [tmp_spec.freq0 if tmp_spec is not None else 1.0]
# Generate model flux
model[d, :, 3] = V0 * (freq/ref_freq)**spi
# Generate gains
jones = None
jones_shape = None
# Dask to NP
t = tbin_idx.compute()
nu = freq.compute()
print('==> Both-mode')
if args.mode == "phase":
jones = phase_gains(lm, nu, n_time, n_ant, args.alpha_std)
elif args.mode == "normal":
jones = normal_gains(t, nu, lm, n_ant, n_corr,
args.sigma_f, args.lt, args.lnu, args.ls)
else:
raise ValueError("Only normal and phase modes available.")
print()
# Reduce jones to diagonals only
jones = jones[:, :, :, :, (0, -1)]
# Jones to complex
jones = jones.astype(np.complex128)
# Jones shape
jones_shape = jones.shape
# Generate filename
if args.out == "":
args.out = f"{args.mode}.npy"
# Save gains and settings to file
with open(args.out, 'wb') as file:
np.save(file, jones)
# Build dask graph
lm = da.from_array(lm, chunks=lm.shape)
model = da.from_array(model, chunks=model.shape)
jones_da = da.from_array(jones, chunks=(args.utimes_per_chunk,)
+ jones_shape[1::])
# Append antenna columns
cols = []
cols.append('ANTENNA1')
cols.append('ANTENNA2')
cols.append('UVW')
# Load data in in chunks and apply gains to each chunk
xds = xds_from_ms(args.ms, columns=cols,
chunks={"row": row_chunks})[0]
ant1 = xds.ANTENNA1.data
ant2 = xds.ANTENNA2.data
# Adjust UVW based on phase-convention
if args.phase_convention == 'CASA':
uvw = -xds.UVW.data.astype(np.float64)
elif args.phase_convention == 'CODEX':
uvw = xds.UVW.data.astype(np.float64)
else:
raise ValueError("Unknown sign convention for phase")
# Get model visibilities
model_vis = np.zeros((n_row, n_freq, n_dir, n_corr),
dtype=np.complex128)
for s in range(n_dir):
model_vis[:, :, s] = im_to_vis(
model[s].reshape((1, n_freq, n_corr)),
uvw,
lm[s].reshape((1, 2)),
freq,
dtype=np.complex64, convention='fourier')
# NP to Dask
model_vis = da.from_array(model_vis, chunks=(row_chunks,
n_freq, n_dir, n_corr))
# Convert Stokes to corr
in_schema = ['I', 'Q', 'U', 'V']
out_schema = [['RR', 'RL'], ['LR', 'LL']]
model_vis = convert(model_vis, in_schema, out_schema)
# Apply gains
data = corrupt_vis(tbin_idx, tbin_counts, ant1, ant2,
jones_da, model_vis).reshape(
(n_row, n_freq, n_corr))
# Assign model visibilities
out_names = []
for d in range(n_dir):
xds = xds.assign(**{source_names[d]:
(("row", "chan", "corr"),
model_vis[:, :, d].reshape(
n_row, n_freq, n_corr).astype(np.complex64))})
out_names += [source_names[d]]
# Assign noise free visibilities to 'CLEAN_DATA'
xds = xds.assign(**{'CLEAN_DATA': (("row", "chan", "corr"),
data.astype(np.complex64))})
out_names += ['CLEAN_DATA']
# Get noise realisation
if args.sigma_n > 0.0:
# Noise matrix
noise = (da.random.normal(loc=0.0, scale=args.sigma_n,
size=(n_row, n_freq, n_corr),
chunks=(row_chunks, n_freq, n_corr)) \
+ 1.0j*da.random.normal(loc=0.0, scale=args.sigma_n,
size=(n_row, n_freq, n_corr),
chunks=(row_chunks, n_freq, n_corr)))/np.sqrt(2.0)
# Zero matrix for off-diagonals
zero = da.zeros_like(noise[:, :, 0])
# Dask to NP
noise = noise.compute()
zero = zero.compute()
# Remove noise on off-diagonals
noise[:, :, 1] = zero[:, :]
noise[:, :, 2] = zero[:, :]
# NP to Dask
noise = da.from_array(noise, chunks=(row_chunks, n_freq, n_corr))
# Assign noise to 'NOISE'
xds = xds.assign(**{'NOISE': (("row", "chan", "corr"),
noise.astype(np.complex64))})
out_names += ['NOISE']
# Add noise to data and assign to 'DATA'
noisy_data = data + noise
xds = xds.assign(**{'DATA': (("row", "chan", "corr"),
noisy_data.astype(np.complex64))})
out_names += ['DATA']
# Create a write to the table
write = xds_to_table(xds, args.ms, out_names)
# Submit all graph computations in parallel
with ProgressBar():
write.compute()
print(f"==> Applied Jones to MS: {args.ms} <--> {args.out}")
def jones(args):
"""Generate jones matrix only, but based off
of a measurement set."""
# Set thread count to cpu count
if args.ncpu:
from multiprocessing.pool import ThreadPool
import dask
dask.config.set(pool=ThreadPool(args.ncpu))
else:
import multiprocessing
args.ncpu = multiprocessing.cpu_count()
# Get full time column and compute row chunks
ms = xds_from_table(args.ms)[0]
_, tbin_idx, tbin_counts = chunkify_rows(
ms.TIME, args.utimes_per_chunk)
# Convert time rows to dask arrays
tbin_idx = da.from_array(tbin_idx,
chunks=(args.utimes_per_chunk))
tbin_counts = da.from_array(tbin_counts,
chunks=(args.utimes_per_chunk))
# Time axis
n_time = tbin_idx.size
# Get antenna columns
ant1 = ms.ANTENNA1.data
ant2 = ms.ANTENNA2.data
# No. of antennas axis
n_ant = (np.maximum(ant1.max(), ant2.max()) + 1).compute()
# Get flag column
flag = ms.FLAG.data
# Get convention
if args.phase_convention == 'CASA':
uvw = -(ms.UVW.data.astype(np.float64))
elif args.phase_convention == 'CODEX':
uvw = ms.UVW.data.astype(np.float64)
else:
raise ValueError("Unknown sign convention for phase")
# Get rest of dimensions
n_row, n_freq, n_corr = flag.shape
# Raise error if correlation axis too small
if n_corr != 4:
raise NotImplementedError("Only 4 correlations "\
+ "currently supported")
# Get phase direction
radec0_table = xds_from_table(args.ms+'::FIELD')[0]
radec0 = radec0_table.PHASE_DIR.data.squeeze().compute()
# Get frequency column
freq_table = xds_from_table(args.ms+'::SPECTRAL_WINDOW')[0]
freq = freq_table.CHAN_FREQ.data.astype(np.float64)[0]
# Check dimension
assert freq.size == n_freq
# Check for sky-model
if args.sky_model == 'MODEL-1.txt':
args.sky_model = MODEL_1
elif args.sky_model == 'MODEL-4.txt':
args.sky_model = MODEL_4
elif args.sky_model == 'MODEL-50.txt':
args.sky_model = MODEL_50
else:
raise ValueError(f"Sky-model {args.sky_model} not in "\
+ "kalcal/datasets/sky_model/")
# Build source model from lsm
lsm = Tigger.load(args.sky_model)
# Direction axis
n_dir = len(lsm.sources)
# Create initial coordinate array and source names
lm = np.zeros((n_dir, 2), dtype=np.float64)
# Cycle coordinates creating a source with flux
for d, source in enumerate(lsm.sources):
# Extract position
radec_s = np.array([[source.pos.ra, source.pos.dec]])
lm[d] = radec_to_lm(radec_s, radec0)
# Generate gains
jones = None
print('==> Jones-only mode')
if args.mode == "phase":
jones = phase_gains(lm, freq, n_time, n_ant, args.alpha_std)
elif args.mode == "normal":
jones = normal_gains(tbin_idx, freq, lm, n_ant, n_corr,
args.sigma_f, args.lt, args.lnu, args.ls)
else:
raise ValueError("Only normal and phase modes available.")
# Reduce jones to diagonals only
jones = jones[:, :, :, :, (0, -1)]
# Jones to complex
jones = jones.astype(np.complex128)
# Generate filename
if args.out == "":
args.out = f"{args.mode}.npy"
# Save gains and settings to file
with open(args.out, 'wb') as file:
np.save(file, jones)
print(f"==> Created Jones data: {args.out}")
def new(ms, sky_model, **kwargs):
"""Generate a jones matrix based on a given sky-model
either as phase-only or normal gains, as an .npy file."""
# Options to attributed dictionary
if kwargs["yaml"] is not None:
options = ocf.load(kwargs["yaml"])
else:
options = ocf.create(kwargs)
# Set to struct
ocf.set_struct(options, True)
# Change path to sky model if chosen
try:
sky_model = sky_models[sky_model.lower()]
except:
# Own sky model reference
pass
# Load ms
MS = xds_from_ms(ms)[0]
# Get dimensions (correlations need to be adapted)
dims = ocf.create(dict(MS.sizes))
n_chan = dims.chan
n_corr = dims.corr
# Get time-bin indices and counts
_, tbin_indices, _ = np.unique(MS.TIME,
return_index=True,
return_counts=True)
# Set time dimension
n_time = len(tbin_indices)
# Get antenna arrays (dask ignored for now)
ant1 = MS.ANTENNA1.data.compute()
ant2 = MS.ANTENNA2.data.compute()
# Set antenna dimension
n_ant = np.max((np.max(ant1), np.max(ant2))) + 1
# Build source model from lsm
lsm = Tigger.load(sky_model)
# Set direction axis as per source
n_dir = len(lsm.sources)
# Get phase direction
radec0_table = xds_from_table(ms + '::FIELD')[0]
radec0 = radec0_table.PHASE_DIR.data.squeeze().compute()
# Get frequency column
freq_table = xds_from_table(ms + '::SPECTRAL_WINDOW')[0]
freq = freq_table.CHAN_FREQ.data.astype(np.float64)[0]
# Check dimension
assert freq.size == n_chan
# Create initial coordinate array and source names
lm = np.zeros((n_dir, 2), dtype=np.float64)
# Cycle coordinates creating a source with flux
for d, source in enumerate(lsm.sources):
# Extract position
radec_s = np.array([[source.pos.ra, source.pos.dec]])
lm[d] = radec_to_lm(radec_s, radec0)
# Direction independent gains
if options.die:
lm = np.array(lm[0]).reshape((1, -1))
n_dir = 1
# Choose between phase-only or normal
if options.type == "phase":
# Run phase-only
print("==> Simulating `phase-only` gains, with dimensions ("\
+ f"n_time={n_time}, n_ant={n_ant}, n_chan={n_chan}, "\
+ f"n_dir={n_dir}, n_corr={n_corr})")
jones = phase_gains(lm, freq, n_time, n_ant, n_chan, n_dir, n_corr,
options.std)
elif options.type == "normal":
# With normal selected, get differentials
lt, lnu, ls = options.diffs
# Run normal
print("==> Simulating `normal` gains, with dimensions ("\
+ f"n_time={n_time}, n_ant={n_ant}, n_chan={n_chan}, "\
+ f"n_dir={n_dir}, n_corr={n_corr})")
jones = normal_gains(tbin_indices, freq, lm, n_time, n_ant, n_chan,
n_dir, n_corr, options.std, lt, lnu, ls)
# Output to jones to .npy file
gains_file = (options.type + ".npy") if options.out_file is None\
else options.out_file
with open(gains_file, 'wb') as file:
np.save(file, jones)
print(f"==> Completed and gains saved to: {gains_file}")
def main(args):
# get full time column and compute row chunks
ms = table(args.ms)
time = ms.getcol('TIME')
row_chunks, tbin_idx, tbin_counts = chunkify_rows(
time, args.utimes_per_chunk)
# convert to dask arrays
tbin_idx = da.from_array(tbin_idx, chunks=(args.utimes_per_chunk))
tbin_counts = da.from_array(tbin_counts, chunks=(args.utimes_per_chunk))
n_time = tbin_idx.size
ms.close()
# get phase dir
fld = table(args.ms+'::FIELD')
radec0 = fld.getcol('PHASE_DIR').squeeze().reshape(1, 2)
radec0 = np.tile(radec0, (n_time, 1))
fld.close()
# get freqs
freqs = table(
args.ms+'::SPECTRAL_WINDOW').getcol('CHAN_FREQ')[0].astype(np.float64)
n_freq = freqs.size
freqs = da.from_array(freqs, chunks=(n_freq))
# get source coordinates from lsm
lsm = Tigger.load(args.sky_model)
radec = []
stokes = []
spi = []
ref_freqs = []
for source in lsm.sources:
radec.append([source.pos.ra, source.pos.dec])
stokes.append([source.flux.I])
spi.append(source.spectrum.spi)
ref_freqs.append(source.spectrum.freq0)
n_dir = len(stokes)
radec = np.asarray(radec)
lm = np.zeros((n_time,) + radec.shape)
for t in range(n_time):
lm[t] = radec_to_lm(radec, radec0[t])
lm = da.from_array(lm, chunks=(args.utimes_per_chunk, n_dir, 2))
# load in the model file
n_corr = 1
model = np.zeros((n_time, n_freq, n_dir, n_corr))
stokes = np.asarray(stokes)
ref_freqs = np.asarray(ref_freqs)
spi = np.asarray(spi)
for t in range(n_time):
for d in range(n_dir):
model[t, :, d, 0] = stokes[d] * (freqs/ref_freqs[d])**spi[d]
# append antenna columns
cols = []
cols.append('ANTENNA1')
cols.append('ANTENNA2')
cols.append('UVW')
# load in gains
jones = np.load(args.gain_file)
jones = jones.astype(np.complex128)
jones_shape = jones.shape
jones = da.from_array(jones, chunks=(args.utimes_per_chunk,)
+ jones_shape[1::])
# change model to dask array
model = da.from_array(model, chunks=(args.utimes_per_chunk,)
+ model.shape[1::])
# load data in in chunks and apply gains to each chunk
xds = xds_from_ms(args.ms, columns=cols, chunks={"row": row_chunks})[0]
ant1 = xds.ANTENNA1.data
ant2 = xds.ANTENNA2.data
uvw = xds.UVW.data
# apply gains
data = compute_and_corrupt_vis(tbin_idx, tbin_counts, ant1, ant2,
jones, model, uvw, freqs, lm)
# Assign visibilities to args.out_col and write to ms
xds = xds.assign(**{args.out_col: (("row", "chan", "corr"), data)})
# Create a write to the table
write = xds_to_table(xds, args.ms, [args.out_col])
# Submit all graph computations in parallel
with ProgressBar():
write.compute()
