def comb_distance(spec_dist, temp_dist, spat_dist):
if logWeight == True:
spec_dist = da.log(spec_dist + 1)
temp_dist = da.log(temp_dist + 1)
comb_dist = da.rechunk(spec_dist * temp_dist * spat_dist,
chunks=spec_dist.chunksize)
print("Done comb distance!", comb_dist)
return comb_dist
def _joint_log_likelihood(self, X):
jll = []
for i in range(np.size(self.classes_)):
jointi = da.log(self.class_prior_[i])
n_ij = -0.5 * da.sum(da.log(2.0 * np.pi * self.sigma_[i, :]))
n_ij -= 0.5 * da.sum(
((X - self.theta_[i, :])**2) / (self.sigma_[i, :]), 1)
jll.append(jointi + n_ij)
joint_log_likelihood = da.stack(jll).T
return joint_log_likelihood
def black_scholes(nopt, price, strike, t, rate, vol, schd=None):
mr = -rate
sig_sig_two = vol * vol * 2
P = price
S = strike
T = t
a = log(P / S)
b = T * mr
z = T * sig_sig_two
c = 0.25 * z
y = da.map_blocks(invsqrt, z)
w1 = (a - b + c) * y
w2 = (a - b - c) * y
d1 = 0.5 + 0.5 * da.map_blocks(erf, w1)
d2 = 0.5 + 0.5 * da.map_blocks(erf, w2)
Se = exp(b) * S
call = P * d1 - Se * d2
put = call - P + Se
return da.compute(da.stack((put, call)), get=schd)
def logsumexp(arr, axis=0):
"""Computes the sum of arr assuming arr is in the log domain.
Returns log(sum(exp(arr))) while minimizing the possibility of
over/underflow.
Examples
--------
>>> import numpy as np
>>> from sklearn.utils.extmath import logsumexp
>>> a = np.arange(10)
>>> np.log(np.sum(np.exp(a)))
9.4586297444267107
>>> logsumexp(a)
9.4586297444267107
"""
if axis == 0:
pass
elif axis == 1:
arr = arr.T
else:
raise NotImplementedError
# Use the max to normalize, as with the log this is what accumulates
# the less errors
vmax = arr.max(axis=0)
out = da.log(da.sum(da.exp(arr - vmax), axis=0))
out += vmax
return out
def score_samples(self, X):
"""Return the log-likelihood of each sample.
See. "Pattern Recognition and Machine Learning"
by C. Bishop, 12.2.1 p. 574
or http://www.miketipping.com/papers/met-mppca.pdf
Parameters
----------
X : array, shape(n_samples, n_features)
The data.
Returns
-------
ll : array, shape (n_samples,)
Log-likelihood of each sample under the current model
"""
check_is_fitted(self, "mean_")
# X = check_array(X)
Xr = X - self.mean_
n_features = X.shape[1]
precision = self.get_precision() # [n_features, n_features]
log_like = -.5 * (Xr * (da.dot(Xr, precision))).sum(axis=1)
log_like -= .5 * (n_features * da.log(2. * np.pi) -
fast_logdet(precision))
return log_like
def black_scholes_numpy_mod(nopt, price, strike, t, rate, vol):
mr = -rate
sig_sig_two = vol * vol * 2
P = price
S = strike
T = t
a = log(P / S)
b = T * mr
z = T * sig_sig_two
c = 0.25 * z
y = invsqrt(z)
w1 = (a - b + c) * y
w2 = (a - b - c) * y
d1 = 0.5 + 0.5 * erf(w1)
d2 = 0.5 + 0.5 * erf(w2)
Se = exp(b) * S
call = P * d1 - Se * d2
put = call - P + Se
return np.stack((call, put))
def _distance(Z, Y, epsilon):
""" Distance function """
Y = Y + epsilon
# The first term below is equal to: da.dot(da.ones(m, n), Y)
# with Z.shape = (m, n) and Y.shape = (n, k)
d = (Y.sum(axis=0, keepdims=True).repeat(Z.shape[0], axis=0) -
da.dot(Z, da.log(Y)))
return d
def _select_by_prediction(self, unlabel_index, predict, batch_size=1):
super()._select_by_prediction(unlabel_index, predict)
entro = []
for vec in predict:
cleanVec = vec.compute()
cleanVec[cleanVec <= 0] = 1e-06 # avoid zero division
entro.append(delayed(sum)(cleanVec * da.log(cleanVec)))
tpl = da.from_array(unlabel_index)
return tpl[nlargestarg(delayed(entro).compute(), batch_size)].compute()
def _U_dask(col):
"""
Compute the U function as part of the thin-plate function.
**Parameter**\n
col: dask.Series or numpy.ndarray
A column vector of ``dask`` (or ``pandas``) ``Series``.
"""
return col**2 * da.log(col.abs())
def get_atm_variables(mus, muv, phi, height, ah2o, bh2o, ao3, tau):
air_mass = 1.0 / mus + 1 / muv
air_mass = air_mass.where(air_mass <= MAXAIRMASS, -1.0)
tO3 = 1.0
tH2O = 1.0
if ao3 != 0:
tO3 = da.exp(-air_mass * UO3 * ao3)
if bh2o != 0:
if bUseV171:
tH2O = da.exp(-da.exp(ah2o + bh2o * da.log(air_mass * UH2O)))
else:
tH2O = da.exp(-(ah2o * ((air_mass * UH2O)**bh2o)))
# Returns sphalb, rhoray, TtotraytH2O, tOG
return atm_variables_finder(mus, muv, phi, height, tau, tO3, tH2O,
TAUSTEP4SPHALB)
def calibration_single_ended_ols(ds, st_label, ast_label):
cal_ref = ds.ufunc_per_section(label=st_label,
ref_temp_broadcasted=True,
calc_per='all')
st = ds.ufunc_per_section(label=st_label, calc_per='all')
ast = ds.ufunc_per_section(label=ast_label, calc_per='all')
z = ds.ufunc_per_section(label='x', calc_per='all')
nx = z.size
nt = ds[st_label].data.shape[1]
p0_est = np.asarray([482., 0.1] + nt * [1.4])
# Eqs for F and B temperature
data1 = 1 / (cal_ref.T.ravel() + 273.15) # gamma
data2 = np.tile(-z, nt) # dalpha
data3 = np.tile([-1.], nt * nx) # C
data = np.concatenate([data1, data2, data3])
# (irow, icol)
coord1row = np.arange(nt * nx, dtype=int)
coord2row = np.arange(nt * nx, dtype=int)
coord3row = np.arange(nt * nx, dtype=int)
coord1col = np.zeros(nt * nx, dtype=int)
coord2col = np.ones(nt * nx, dtype=int)
coord3col = np.repeat(np.arange(2, nt + 2, dtype=int), nx)
rows = [coord1row, coord2row, coord3row]
cols = [coord1col, coord2col, coord3col]
coords = (np.concatenate(rows), np.concatenate(cols))
# try scipy.sparse.bsr_matrix
X = sp.coo_matrix((data, coords),
shape=(nt * nx, nt + 2),
dtype=float,
copy=False)
y = da.log(st / ast).T.ravel()
# noinspection PyTypeChecker
p0 = ln.lsqr(X, y, x0=p0_est, show=True, calc_var=True)
return nt, z, p0
def _ir_calibrate(self, data):
"""IR calibration."""
cwl = self._header['block5']["central_wave_length"][0] * 1e-6
c__ = self._header['calibration']["speed_of_light"][0]
h__ = self._header['calibration']["planck_constant"][0]
k__ = self._header['calibration']["boltzmann_constant"][0]
a__ = (h__ * c__) / (k__ * cwl)
b__ = ((2 * h__ * c__**2) / (data * 1.0e6 * cwl**5)) + 1
Te_ = a__ / da.log(b__)
c0_ = self._header['calibration']["c0_rad2tb_conversion"][0]
c1_ = self._header['calibration']["c1_rad2tb_conversion"][0]
c2_ = self._header['calibration']["c2_rad2tb_conversion"][0]
return (c0_ + c1_ * Te_ + c2_ * Te_**2).clip(0)
def _loss(self, prob):
"""Compute expected log-loss.
Parameters
----------
prob: 2d array, shape [n_samples, n_classes]
The probabilistic prediction matrix for the unlabeled set.
Returns
-------
log_loss: float
The sum of log_loss for the prob.
"""
log_loss = []
for i in range(len(prob)):
for p in list(prob[i]):
log_loss.append(delayed(p * da.log(p)))
return delayed(sum)(log_loss).compute()
def _ir_calibrate(self, data):
"""IR calibration
"""
cwl = self._header['block5']["central_wave_length"][0] * 1e-6
c__ = self._header['calibration']["speed_of_light"][0]
h__ = self._header['calibration']["planck_constant"][0]
k__ = self._header['calibration']["boltzmann_constant"][0]
a__ = (h__ * c__) / (k__ * cwl)
b__ = ((2 * h__ * c__ ** 2) / (data * 1.0e6 * cwl ** 5)) + 1
Te_ = a__ / da.log(b__)
c0_ = self._header['calibration']["c0_rad2tb_conversion"][0]
c1_ = self._header['calibration']["c1_rad2tb_conversion"][0]
c2_ = self._header['calibration']["c2_rad2tb_conversion"][0]
return (c0_ + c1_ * Te_ + c2_ * Te_ ** 2).clip(0)
def predict_log_proba(self, X):
"""Log of probability estimates.
For dask inputs, a dask array or dataframe is returned. For other
inputs (NumPy array, pandas dataframe, scipy sparse matrix), the
regular return value is returned.
If the underlying estimator does not have a ``predict_proba``
method, then an ``AttributeError`` is raised.
Parameters
----------
X : array or dataframe
Returns
-------
y : array-like
"""
self._check_method("predict_log_proba")
return da.log(self.predict_proba(X))
def test_arithmetic():
x = np.arange(5).astype('f4') + 2
y = np.arange(5).astype('i8') + 2
z = np.arange(5).astype('i4') + 2
a = da.from_array(x, chunks=(2,))
b = da.from_array(y, chunks=(2,))
c = da.from_array(z, chunks=(2,))
assert eq(a + b, x + y)
assert eq(a * b, x * y)
assert eq(a - b, x - y)
assert eq(a / b, x / y)
assert eq(b & b, y & y)
assert eq(b | b, y | y)
assert eq(b ^ b, y ^ y)
assert eq(a // b, x // y)
assert eq(a ** b, x ** y)
assert eq(a % b, x % y)
assert eq(a > b, x > y)
assert eq(a < b, x < y)
assert eq(a >= b, x >= y)
assert eq(a <= b, x <= y)
assert eq(a == b, x == y)
assert eq(a != b, x != y)
assert eq(a + 2, x + 2)
assert eq(a * 2, x * 2)
assert eq(a - 2, x - 2)
assert eq(a / 2, x / 2)
assert eq(b & True, y & True)
assert eq(b | True, y | True)
assert eq(b ^ True, y ^ True)
assert eq(a // 2, x // 2)
assert eq(a ** 2, x ** 2)
assert eq(a % 2, x % 2)
assert eq(a > 2, x > 2)
assert eq(a < 2, x < 2)
assert eq(a >= 2, x >= 2)
assert eq(a <= 2, x <= 2)
assert eq(a == 2, x == 2)
assert eq(a != 2, x != 2)
assert eq(2 + b, 2 + y)
assert eq(2 * b, 2 * y)
assert eq(2 - b, 2 - y)
assert eq(2 / b, 2 / y)
assert eq(True & b, True & y)
assert eq(True | b, True | y)
assert eq(True ^ b, True ^ y)
assert eq(2 // b, 2 // y)
assert eq(2 ** b, 2 ** y)
assert eq(2 % b, 2 % y)
assert eq(2 > b, 2 > y)
assert eq(2 < b, 2 < y)
assert eq(2 >= b, 2 >= y)
assert eq(2 <= b, 2 <= y)
assert eq(2 == b, 2 == y)
assert eq(2 != b, 2 != y)
assert eq(-a, -x)
assert eq(abs(a), abs(x))
assert eq(~(a == b), ~(x == y))
assert eq(~(a == b), ~(x == y))
assert eq(da.logaddexp(a, b), np.logaddexp(x, y))
assert eq(da.logaddexp2(a, b), np.logaddexp2(x, y))
assert eq(da.exp(b), np.exp(y))
assert eq(da.log(a), np.log(x))
assert eq(da.log10(a), np.log10(x))
assert eq(da.log1p(a), np.log1p(x))
assert eq(da.expm1(b), np.expm1(y))
assert eq(da.sqrt(a), np.sqrt(x))
assert eq(da.square(a), np.square(x))
assert eq(da.sin(a), np.sin(x))
assert eq(da.cos(b), np.cos(y))
assert eq(da.tan(a), np.tan(x))
assert eq(da.arcsin(b/10), np.arcsin(y/10))
assert eq(da.arccos(b/10), np.arccos(y/10))
assert eq(da.arctan(b/10), np.arctan(y/10))
assert eq(da.arctan2(b*10, a), np.arctan2(y*10, x))
assert eq(da.hypot(b, a), np.hypot(y, x))
assert eq(da.sinh(a), np.sinh(x))
assert eq(da.cosh(b), np.cosh(y))
assert eq(da.tanh(a), np.tanh(x))
assert eq(da.arcsinh(b*10), np.arcsinh(y*10))
assert eq(da.arccosh(b*10), np.arccosh(y*10))
assert eq(da.arctanh(b/10), np.arctanh(y/10))
assert eq(da.deg2rad(a), np.deg2rad(x))
assert eq(da.rad2deg(a), np.rad2deg(x))
assert eq(da.logical_and(a < 1, b < 4), np.logical_and(x < 1, y < 4))
assert eq(da.logical_or(a < 1, b < 4), np.logical_or(x < 1, y < 4))
assert eq(da.logical_xor(a < 1, b < 4), np.logical_xor(x < 1, y < 4))
assert eq(da.logical_not(a < 1), np.logical_not(x < 1))
assert eq(da.maximum(a, 5 - a), np.maximum(a, 5 - a))
assert eq(da.minimum(a, 5 - a), np.minimum(a, 5 - a))
assert eq(da.fmax(a, 5 - a), np.fmax(a, 5 - a))
assert eq(da.fmin(a, 5 - a), np.fmin(a, 5 - a))
assert eq(da.isreal(a + 1j * b), np.isreal(x + 1j * y))
assert eq(da.iscomplex(a + 1j * b), np.iscomplex(x + 1j * y))
assert eq(da.isfinite(a), np.isfinite(x))
assert eq(da.isinf(a), np.isinf(x))
assert eq(da.isnan(a), np.isnan(x))
assert eq(da.signbit(a - 3), np.signbit(x - 3))
assert eq(da.copysign(a - 3, b), np.copysign(x - 3, y))
assert eq(da.nextafter(a - 3, b), np.nextafter(x - 3, y))
assert eq(da.ldexp(c, c), np.ldexp(z, z))
assert eq(da.fmod(a * 12, b), np.fmod(x * 12, y))
assert eq(da.floor(a * 0.5), np.floor(x * 0.5))
assert eq(da.ceil(a), np.ceil(x))
assert eq(da.trunc(a / 2), np.trunc(x / 2))
assert eq(da.degrees(b), np.degrees(y))
assert eq(da.radians(a), np.radians(x))
assert eq(da.rint(a + 0.3), np.rint(x + 0.3))
assert eq(da.fix(a - 2.5), np.fix(x - 2.5))
assert eq(da.angle(a + 1j), np.angle(x + 1j))
assert eq(da.real(a + 1j), np.real(x + 1j))
assert eq((a + 1j).real, np.real(x + 1j))
assert eq(da.imag(a + 1j), np.imag(x + 1j))
assert eq((a + 1j).imag, np.imag(x + 1j))
assert eq(da.conj(a + 1j * b), np.conj(x + 1j * y))
assert eq((a + 1j * b).conj(), (x + 1j * y).conj())
assert eq(da.clip(b, 1, 4), np.clip(y, 1, 4))
assert eq(da.fabs(b), np.fabs(y))
assert eq(da.sign(b - 2), np.sign(y - 2))
l1, l2 = da.frexp(a)
r1, r2 = np.frexp(x)
assert eq(l1, r1)
assert eq(l2, r2)
l1, l2 = da.modf(a)
r1, r2 = np.modf(x)
assert eq(l1, r1)
assert eq(l2, r2)
assert eq(da.around(a, -1), np.around(x, -1))
def radiance_to_bt(arr, wc_, a__, b__):
"""Convert to BT.
"""
return a__ + b__ * (C2 * wc_ / (da.log(1 + (C1 * (wc_**3) / arr))))
def _predict(args):
# get inclusion regions
include_regions = load_regions(args.within) if args.within else []
# Import source data from WSClean component list
# See https://sourceforge.net/p/wsclean/wiki/ComponentList
(comp_type, radec, stokes, spec_coeff, ref_freq, log_spec_ind,
gaussian_shape) = import_from_wsclean(args.sky_model,
include_regions=include_regions,
point_only=args.points_only,
num=args.num_sources or None)
# Add output column if it isn't present
ms_rows, ms_datatype = ms_preprocess(args)
# Get the support tables
tables = support_tables(
args, ["FIELD", "DATA_DESCRIPTION", "SPECTRAL_WINDOW", "POLARIZATION"])
field_ds = tables["FIELD"]
ddid_ds = tables["DATA_DESCRIPTION"]
spw_ds = tables["SPECTRAL_WINDOW"]
pol_ds = tables["POLARIZATION"]
max_num_chan = max([ss.NUM_CHAN.data[0] for ss in spw_ds])
max_num_corr = max([ss.NUM_CORR.data[0] for ss in pol_ds])
# Perform resource budgeting
args.row_chunks, args.model_chunks = get_budget(comp_type.shape[0],
ms_rows, max_num_chan,
max_num_corr, ms_datatype,
args)
radec = da.from_array(radec, chunks=(args.model_chunks, 2))
stokes = da.from_array(stokes, chunks=(args.model_chunks, 4))
if np.count_nonzero(comp_type == 'GAUSSIAN') > 0:
gaussian_components = True
gshape_chunks = (args.model_chunks, 3)
gaussian_shape = da.from_array(gaussian_shape, chunks=gshape_chunks)
else:
gaussian_components = False
if args.spectra:
spec_chunks = (args.model_chunks, spec_coeff.shape[1])
spec_coeff = da.from_array(spec_coeff, chunks=spec_chunks)
ref_freq = da.from_array(ref_freq, chunks=(args.model_chunks, ))
# List of write operations
writes = []
# Construct a graph for each FIELD and DATA DESCRIPTOR
datasets = xds_from_ms(args.ms,
columns=["UVW", "ANTENNA1", "ANTENNA2", "TIME"],
group_cols=["FIELD_ID", "DATA_DESC_ID"],
chunks={"row": args.row_chunks})
select_fields = valid_field_ids(field_ds, args.fields)
for xds in filter_datasets(datasets, select_fields):
# Extract frequencies from the spectral window associated
# with this data descriptor id
field = field_ds[xds.attrs['FIELD_ID']]
ddid = ddid_ds[xds.attrs['DATA_DESC_ID']]
spw = spw_ds[ddid.SPECTRAL_WINDOW_ID.data[0]]
pol = pol_ds[ddid.POLARIZATION_ID.data[0]]
frequency = spw.CHAN_FREQ.data[0]
corrs = pol.NUM_CORR.values
lm = radec_to_lm(radec, field.PHASE_DIR.data[0][0])
if args.exp_sign_convention == 'casa':
uvw = -xds.UVW.data
elif args.exp_sign_convention == 'thompson':
uvw = xds.UVW.data
else:
raise ValueError("Invalid sign convention '%s'" % args.sign)
if args.spectra:
# flux density at reference frequency ...
# ... for logarithmic polynomial functions
if log_spec_ind:
Is = da.log(stokes[:, 0, None]) * frequency[None, :]**0
# ... or for ordinary polynomial functions
else:
Is = stokes[:, 0, None] * frequency[None, :]**0
# additional terms of SED ...
for jj in range(spec_coeff.shape[1]):
# ... for logarithmic polynomial functions
if log_spec_ind:
Is += spec_coeff[:, jj, None] * \
da.log((frequency[None, :]/ref_freq[:, None])**(jj+1))
# ... or for ordinary polynomial functions
else:
Is += spec_coeff[:, jj, None] * \
(frequency[None, :]/ref_freq[:, None]-1)**(jj+1)
if log_spec_ind:
Is = da.exp(Is)
Qs = da.zeros_like(Is)
Us = da.zeros_like(Is)
Vs = da.zeros_like(Is)
# stack along new axis and make it the last axis of the new array
spectrum = da.stack([Is, Qs, Us, Vs], axis=-1)
spectrum = spectrum.rechunk(spectrum.chunks[:2] +
(spectrum.shape[2], ))
print('-------------------------------------------')
print('Nr sources = {0:d}'.format(stokes.shape[0]))
print('-------------------------------------------')
print('stokes.shape = {0:}'.format(stokes.shape))
print('frequency.shape = {0:}'.format(frequency.shape))
if args.spectra:
print('Is.shape = {0:}'.format(Is.shape))
if args.spectra:
print('spectrum.shape = {0:}'.format(spectrum.shape))
# (source, row, frequency)
phase = phase_delay(lm, uvw, frequency)
# If at least one Gaussian component is present in the component
# list then all sources are modelled as Gaussian components
# (Delta components have zero width)
if gaussian_components:
phase *= gaussian(uvw, frequency, gaussian_shape)
# (source, frequency, corr_products)
brightness = convert(spectrum if args.spectra else stokes,
["I", "Q", "U", "V"], corr_schema(pol))
print('brightness.shape = {0:}'.format(brightness.shape))
print('phase.shape = {0:}'.format(phase.shape))
print('-------------------------------------------')
print('Attempting phase-brightness einsum with "{0:s}"'.format(
einsum_schema(pol, args.spectra)))
# (source, row, frequency, corr_products)
jones = da.einsum(einsum_schema(pol, args.spectra), phase, brightness)
print('jones.shape = {0:}'.format(jones.shape))
print('-------------------------------------------')
if gaussian_components:
print('Some Gaussian sources found')
else:
print('All sources are Delta functions')
print('-------------------------------------------')
# Identify time indices
_, time_index = da.unique(xds.TIME.data, return_inverse=True)
# Predict visibilities
vis = predict_vis(time_index, xds.ANTENNA1.data, xds.ANTENNA2.data,
None, jones, None, None, None, None)
# Reshape (2, 2) correlation to shape (4,)
if corrs == 4:
vis = vis.reshape(vis.shape[:2] + (4, ))
# Assign visibilities to MODEL_DATA array on the dataset
xds = xds.assign(
**{args.output_column: (("row", "chan", "corr"), vis)})
# Create a write to the table
write = xds_to_table(xds, args.ms, [args.output_column])
# Add to the list of writes
writes.append(write)
with ExitStack() as stack:
if sys.stdout.isatty():
# Default progress bar in user terminal
stack.enter_context(ProgressBar())
else:
# Log progress every 5 minutes
stack.enter_context(ProgressBar(minimum=2 * 60, dt=5))
# Submit all graph computations in parallel
dask.compute(writes)
def _distance(Z, Y, epsilon):
""" Distance function """
Y = Y + epsilon
# The first term below is equal to one row of: da.dot(da.ones(m, n), Y)
# with Z.shape = (m, n) and Y.shape = (n, k)
return Y.sum(axis=0, keepdims=True) - da.matmul(Z, da.log(Y))
def _band_log(arr):
slope = (factor - 1.) / float(arr.max() - arr.min())
arr = 1. + (arr - arr.min()) * slope
arr = c__ + b__ * da.log(arr)
return arr
def calibration_double_ended_wls(ds,
st_label,
ast_label,
rst_label,
rast_label,
st_var,
ast_var,
rst_var,
rast_var,
calc_cov=True,
solver='sparse',
dtype32=False):
"""
Parameters
----------
ds : DataStore
st_label
ast_label
rst_label
rast_label
st_var
ast_var
rst_var
rast_var
calc_cov
solver : {'sparse', 'stats'}
Returns
-------
"""
# x_alpha_set_zero=0.,
# set one alpha for all times to zero
# x_alpha_set_zeroi = np.argmin(np.abs(ds.x.data - x_alpha_set_zero))
# x_alpha_set_zeroidata = np.arange(nt) * no + x_alpha_set_zeroi
cal_ref = ds.ufunc_per_section(label=st_label,
ref_temp_broadcasted=True,
calc_per='all')
st = ds.ufunc_per_section(label=st_label, calc_per='all')
ast = ds.ufunc_per_section(label=ast_label, calc_per='all')
rst = ds.ufunc_per_section(label=rst_label, calc_per='all')
rast = ds.ufunc_per_section(label=rast_label, calc_per='all')
z = ds.ufunc_per_section(label='x', calc_per='all')
nx = z.size
_xsorted = np.argsort(ds.x.data)
_ypos = np.searchsorted(ds.x.data[_xsorted], z)
x_index = _xsorted[_ypos]
no, nt = ds[st_label].data.shape
p0_est = np.asarray([482., 0.1] + nt * [1.4] + no * [0.])
# Data for F and B temperature, 2 * nt * nx items
data1 = da.repeat(1 / (cal_ref.T.ravel() + 273.15), 2) # gamma
# data2 = da.tile(np.array([0., -1.]), nt * nx) # alphaint
data2 = da.stack((da.zeros(nt * nx, chunks=nt * nx),
-da.ones(nt * nx, chunks=nt * nx))).T.ravel()
# data3 = da.tile(np.array([-1., -1.]), nt * nx) # C
data3 = -da.ones(2 * nt * nx, chunks=2 * nt * nx)
# data5 = da.tile(np.array([-1., 1.]), nt * nx) # alph
data5 = da.stack((-da.ones(nt * nx, chunks=nt * nx),
da.ones(nt * nx, chunks=nt * nx))).T.ravel()
# Data for alpha, nt * no items
# data6 = da.repeat(np.array([-0.5]), nt * no) # alphaint
data6 = da.ones(nt * no, dtype=float,
chunks=(nt * no, )) * -0.5 # alphaint
data9 = da.ones(nt * no, dtype=float, chunks=(nt * no, )) # alpha
# alpha should start at zero. But then the sparse solver crashes
# data9[x_alpha_set_zeroidata] = 0.
data = da.concatenate([data1, data2, data3, data5, data6, data9]).compute()
# Coords (irow, icol)
coord1row = da.arange(2 * nt * nx, dtype=int, chunks=(nt * nx, )) # gamma
coord2row = da.arange(2 * nt * nx, dtype=int,
chunks=(nt * nx, )) # alphaint
coord3row = da.arange(2 * nt * nx, dtype=int, chunks=(nt * nx, )) # C
coord5row = da.arange(2 * nt * nx, dtype=int, chunks=(nt * nx, )) # alpha
coord6row = da.arange(2 * nt * nx,
2 * nt * nx + nt * no,
dtype=int,
chunks=(nt * no, )) # alphaint
coord9row = da.arange(2 * nt * nx,
2 * nt * nx + nt * no,
dtype=int,
chunks=(nt * no, )) # alpha
coord1col = da.zeros(2 * nt * nx, dtype=int, chunks=(nt * nx, )) # gamma
coord2col = da.ones(2 * nt * nx, dtype=int, chunks=(nt * nx, )) * (
2 + nt + no - 1) # alphaint
coord3col = da.repeat(da.arange(nt, dtype=int, chunks=(nt, )) + 2,
2 * nx).rechunk(nt * nx) # C
coord5col = da.tile(np.repeat(x_index, 2) + nt + 2,
nt).rechunk(nt * nx) # alpha
coord6col = da.ones(nt * no, dtype=int,
chunks=(nt * no, )) # * (2 + nt + no - 1) # alphaint
coord9col = da.tile(
da.arange(no, dtype=int, chunks=(nt * no, )) + nt + 2, nt) # alpha
rows = [coord1row, coord2row, coord3row, coord5row, coord6row, coord9row]
cols = [coord1col, coord2col, coord3col, coord5col, coord6col, coord9col]
coords = (da.concatenate(rows).compute(), da.concatenate(cols).compute())
# try scipy.sparse.bsr_matrix
X = sp.coo_matrix((data, coords),
shape=(2 * nx * nt + nt * no, nt + 2 + no),
dtype=float,
copy=False)
# Spooky way to interleave and ravel arrays in correct order. Works!
y1F = da.log(st / ast).T.ravel()
y1B = da.log(rst / rast).T.ravel()
y1 = da.stack([y1F, y1B]).T.ravel()
y2F = da.log(ds[st_label].data / ds[ast_label].data).T.ravel()
y2B = da.log(ds[rst_label].data / ds[rast_label].data).T.ravel()
y2 = (y2B - y2F) / 2
y = da.concatenate([y1, y2]).compute()
# Calculate the reprocical of the variance (not std)
w1F = (1 / st**2 * st_var + 1 / ast**2 * ast_var).T.ravel()
w1B = (1 / rst**2 * rst_var + 1 / rast**2 * rast_var).T.ravel()
w1 = da.stack([w1F, w1B]).T.ravel()
w2 = (0.5 / ds[st_label].data**2 * st_var +
0.5 / ds[ast_label].data**2 * ast_var +
0.5 / ds[rst_label].data**2 * rst_var +
0.5 / ds[rast_label].data**2 * rast_var).T.ravel()
w = da.concatenate([w1, w2]).compute()
if solver == 'sparse':
p_sol, p_var, p_cov = wls_sparse(X,
y,
w=w,
x0=p0_est,
calc_cov=calc_cov,
dtype32=dtype32)
elif solver == 'stats':
p_sol, p_var, p_cov = wls_stats(X, y, w=w, calc_cov=calc_cov)
if calc_cov:
return nt, z, p_sol, p_var, p_cov
else:
return nt, z, p_sol, p_var
def calibration_single_ended_wls(ds,
st_label,
ast_label,
st_var,
ast_var,
calc_cov=True,
solver='sparse'):
cal_ref = ds.ufunc_per_section(label=st_label,
ref_temp_broadcasted=True,
calc_per='all')
st = ds.ufunc_per_section(label=st_label, calc_per='all')
ast = ds.ufunc_per_section(label=ast_label, calc_per='all')
z = ds.ufunc_per_section(label='x', calc_per='all')
nx = z.size
nt = ds[st_label].data.shape[1]
p0_est = np.asarray([482., 0.1] + nt * [1.4])
# Eqs for F and B temperature
data1 = 1 / (cal_ref.T.ravel() + 273.15) # gamma
data2 = np.tile(-z, nt) # dalpha
data3 = np.tile([-1.], nt * nx) # C
data = np.concatenate([data1, data2, data3])
# (irow, icol)
coord1row = np.arange(nt * nx, dtype=int)
coord2row = np.arange(nt * nx, dtype=int)
coord3row = np.arange(nt * nx, dtype=int)
coord1col = np.zeros(nt * nx, dtype=int)
coord2col = np.ones(nt * nx, dtype=int)
coord3col = np.repeat(np.arange(2, nt + 2, dtype=int), nx)
rows = [coord1row, coord2row, coord3row]
cols = [coord1col, coord2col, coord3col]
coords = (np.concatenate(rows), np.concatenate(cols))
# try scipy.sparse.bsr_matrix
X = sp.coo_matrix((data, coords),
shape=(nt * nx, nt + 2),
dtype=float,
copy=False)
y = da.log(st / ast).T.ravel().compute()
w = (1 / st**2 * st_var + 1 / ast**2 * ast_var).T.ravel().compute()
if solver == 'sparse':
p_sol, p_var, p_cov = wls_sparse(X,
y,
w=w,
x0=p0_est,
calc_cov=calc_cov)
elif solver == 'stats':
p_sol, p_var, p_cov = wls_stats(X, y, w=w, calc_cov=calc_cov)
else:
raise ValueError("Choose a valid solver")
if calc_cov:
return nt, z, p_sol, p_var, p_cov
else:
return nt, z, p_sol, p_var
def _band_log(arr):
slope = (factor - 1.) / float(arr.max() - arr.min())
arr = 1. + (arr - arr.min()) * slope
arr = c__ + b__ * da.log(arr)
return arr
def calibration_double_ended_ols(ds, st_label, ast_label, rst_label,
rast_label):
cal_ref = ds.ufunc_per_section(label=st_label,
ref_temp_broadcasted=True,
calc_per='all')
st = ds.ufunc_per_section(label=st_label, calc_per='all')
ast = ds.ufunc_per_section(label=ast_label, calc_per='all')
rst = ds.ufunc_per_section(label=rst_label, calc_per='all')
rast = ds.ufunc_per_section(label=rast_label, calc_per='all')
z = ds.ufunc_per_section(label='x', calc_per='all')
nx = z.size
_xsorted = np.argsort(ds.x.data)
_ypos = np.searchsorted(ds.x.data[_xsorted], z)
x_index = _xsorted[_ypos]
# if hasattr(cal_ref, 'chunks'):
# chunks_dim = (nx, cal_ref.chunks[1])
#
# for item in [st, ast, rst, rast]:
# item.rechunk(chunks_dim)
nt = ds[st_label].data.shape[1]
no = ds[st_label].data.shape[0]
p0_est = np.asarray([482., 0.1] + nt * [1.4] + no * [0.])
# Eqs for F and B temperature
data1 = np.repeat(1 / (cal_ref.T.ravel() + 273.15), 2) # gamma
data2 = np.tile([0., -1.], nt * nx) # alphaint
data3 = np.tile([-1., -1.], nt * nx) # C
data5 = np.tile([-1., 1.], nt * nx) # alpha
# Eqs for alpha
data6 = np.repeat(-0.5, nt * no) # alphaint
data9 = np.ones(nt * no, dtype=float) # alpha
data = np.concatenate([data1, data2, data3, data5, data6, data9])
# (irow, icol)
coord1row = np.arange(2 * nt * nx, dtype=int)
coord2row = np.arange(2 * nt * nx, dtype=int)
coord3row = np.arange(2 * nt * nx, dtype=int)
coord5row = np.arange(2 * nt * nx, dtype=int)
coord6row = np.arange(2 * nt * nx, 2 * nt * nx + nt * no, dtype=int)
coord9row = np.arange(2 * nt * nx, 2 * nt * nx + nt * no, dtype=int)
coord1col = np.zeros(2 * nt * nx, dtype=int)
coord2col = np.ones(2 * nt * nx, dtype=int) * (2 + nt + no - 1)
coord3col = np.repeat(np.arange(nt, dtype=int) + 2, 2 * nx)
coord5col = np.tile(np.repeat(x_index, 2) + nt + 2, nt)
coord6col = np.ones(nt * no, dtype=int)
coord9col = np.tile(np.arange(no, dtype=int) + nt + 2, nt)
rows = [coord1row, coord2row, coord3row, coord5row, coord6row, coord9row]
cols = [coord1col, coord2col, coord3col, coord5col, coord6col, coord9col]
coords = (np.concatenate(rows), np.concatenate(cols))
# try scipy.sparse.bsr_matrix
X = sp.coo_matrix((data, coords),
shape=(2 * nx * nt + nt * no, nt + 2 + no),
dtype=float,
copy=False)
y1F = da.log(st / ast).T.ravel()
y1B = da.log(rst / rast).T.ravel()
y1 = da.stack([y1F, y1B]).T.ravel()
y2F = np.log(ds[st_label].data / ds[ast_label].data).T.ravel()
y2B = np.log(ds[rst_label].data / ds[rast_label].data).T.ravel()
y2 = (y2B - y2F) / 2
y = da.concatenate([y1, y2]).compute()
# noinspection PyTypeChecker
p0 = ln.lsqr(X, y, x0=p0_est, show=True, calc_var=True)
return nt, z, p0
def test_arithmetic():
x = np.arange(5).astype('f4') + 2
y = np.arange(5).astype('i8') + 2
z = np.arange(5).astype('i4') + 2
a = da.from_array(x, chunks=(2, ))
b = da.from_array(y, chunks=(2, ))
c = da.from_array(z, chunks=(2, ))
assert eq(a + b, x + y)
assert eq(a * b, x * y)
assert eq(a - b, x - y)
assert eq(a / b, x / y)
assert eq(b & b, y & y)
assert eq(b | b, y | y)
assert eq(b ^ b, y ^ y)
assert eq(a // b, x // y)
assert eq(a**b, x**y)
assert eq(a % b, x % y)
assert eq(a > b, x > y)
assert eq(a < b, x < y)
assert eq(a >= b, x >= y)
assert eq(a <= b, x <= y)
assert eq(a == b, x == y)
assert eq(a != b, x != y)
assert eq(a + 2, x + 2)
assert eq(a * 2, x * 2)
assert eq(a - 2, x - 2)
assert eq(a / 2, x / 2)
assert eq(b & True, y & True)
assert eq(b | True, y | True)
assert eq(b ^ True, y ^ True)
assert eq(a // 2, x // 2)
assert eq(a**2, x**2)
assert eq(a % 2, x % 2)
assert eq(a > 2, x > 2)
assert eq(a < 2, x < 2)
assert eq(a >= 2, x >= 2)
assert eq(a <= 2, x <= 2)
assert eq(a == 2, x == 2)
assert eq(a != 2, x != 2)
assert eq(2 + b, 2 + y)
assert eq(2 * b, 2 * y)
assert eq(2 - b, 2 - y)
assert eq(2 / b, 2 / y)
assert eq(True & b, True & y)
assert eq(True | b, True | y)
assert eq(True ^ b, True ^ y)
assert eq(2 // b, 2 // y)
assert eq(2**b, 2**y)
assert eq(2 % b, 2 % y)
assert eq(2 > b, 2 > y)
assert eq(2 < b, 2 < y)
assert eq(2 >= b, 2 >= y)
assert eq(2 <= b, 2 <= y)
assert eq(2 == b, 2 == y)
assert eq(2 != b, 2 != y)
assert eq(-a, -x)
assert eq(abs(a), abs(x))
assert eq(~(a == b), ~(x == y))
assert eq(~(a == b), ~(x == y))
assert eq(da.logaddexp(a, b), np.logaddexp(x, y))
assert eq(da.logaddexp2(a, b), np.logaddexp2(x, y))
assert eq(da.exp(b), np.exp(y))
assert eq(da.log(a), np.log(x))
assert eq(da.log10(a), np.log10(x))
assert eq(da.log1p(a), np.log1p(x))
assert eq(da.expm1(b), np.expm1(y))
assert eq(da.sqrt(a), np.sqrt(x))
assert eq(da.square(a), np.square(x))
assert eq(da.sin(a), np.sin(x))
assert eq(da.cos(b), np.cos(y))
assert eq(da.tan(a), np.tan(x))
assert eq(da.arcsin(b / 10), np.arcsin(y / 10))
assert eq(da.arccos(b / 10), np.arccos(y / 10))
assert eq(da.arctan(b / 10), np.arctan(y / 10))
assert eq(da.arctan2(b * 10, a), np.arctan2(y * 10, x))
assert eq(da.hypot(b, a), np.hypot(y, x))
assert eq(da.sinh(a), np.sinh(x))
assert eq(da.cosh(b), np.cosh(y))
assert eq(da.tanh(a), np.tanh(x))
assert eq(da.arcsinh(b * 10), np.arcsinh(y * 10))
assert eq(da.arccosh(b * 10), np.arccosh(y * 10))
assert eq(da.arctanh(b / 10), np.arctanh(y / 10))
assert eq(da.deg2rad(a), np.deg2rad(x))
assert eq(da.rad2deg(a), np.rad2deg(x))
assert eq(da.logical_and(a < 1, b < 4), np.logical_and(x < 1, y < 4))
assert eq(da.logical_or(a < 1, b < 4), np.logical_or(x < 1, y < 4))
assert eq(da.logical_xor(a < 1, b < 4), np.logical_xor(x < 1, y < 4))
assert eq(da.logical_not(a < 1), np.logical_not(x < 1))
assert eq(da.maximum(a, 5 - a), np.maximum(a, 5 - a))
assert eq(da.minimum(a, 5 - a), np.minimum(a, 5 - a))
assert eq(da.fmax(a, 5 - a), np.fmax(a, 5 - a))
assert eq(da.fmin(a, 5 - a), np.fmin(a, 5 - a))
assert eq(da.isreal(a + 1j * b), np.isreal(x + 1j * y))
assert eq(da.iscomplex(a + 1j * b), np.iscomplex(x + 1j * y))
assert eq(da.isfinite(a), np.isfinite(x))
assert eq(da.isinf(a), np.isinf(x))
assert eq(da.isnan(a), np.isnan(x))
assert eq(da.signbit(a - 3), np.signbit(x - 3))
assert eq(da.copysign(a - 3, b), np.copysign(x - 3, y))
assert eq(da.nextafter(a - 3, b), np.nextafter(x - 3, y))
assert eq(da.ldexp(c, c), np.ldexp(z, z))
assert eq(da.fmod(a * 12, b), np.fmod(x * 12, y))
assert eq(da.floor(a * 0.5), np.floor(x * 0.5))
assert eq(da.ceil(a), np.ceil(x))
assert eq(da.trunc(a / 2), np.trunc(x / 2))
assert eq(da.degrees(b), np.degrees(y))
assert eq(da.radians(a), np.radians(x))
assert eq(da.rint(a + 0.3), np.rint(x + 0.3))
assert eq(da.fix(a - 2.5), np.fix(x - 2.5))
assert eq(da.angle(a + 1j), np.angle(x + 1j))
assert eq(da.real(a + 1j), np.real(x + 1j))
assert eq((a + 1j).real, np.real(x + 1j))
assert eq(da.imag(a + 1j), np.imag(x + 1j))
assert eq((a + 1j).imag, np.imag(x + 1j))
assert eq(da.conj(a + 1j * b), np.conj(x + 1j * y))
assert eq((a + 1j * b).conj(), (x + 1j * y).conj())
assert eq(da.clip(b, 1, 4), np.clip(y, 1, 4))
assert eq(da.fabs(b), np.fabs(y))
assert eq(da.sign(b - 2), np.sign(y - 2))
l1, l2 = da.frexp(a)
r1, r2 = np.frexp(x)
assert eq(l1, r1)
assert eq(l2, r2)
l1, l2 = da.modf(a)
r1, r2 = np.modf(x)
assert eq(l1, r1)
assert eq(l2, r2)
assert eq(da.around(a, -1), np.around(x, -1))
def chand(phi, muv, mus, taur):
# FROM FUNCTION CHAND
# phi: azimuthal difference between sun and observation in degree
# (phi=0 in backscattering direction)
# mus: cosine of the sun zenith angle
# muv: cosine of the observation zenith angle
# taur: molecular optical depth
# rhoray: molecular path reflectance
# constant xdep: depolarization factor (0.0279)
# xfd = (1-xdep/(2-xdep)) / (1 + 2*xdep/(2-xdep)) = 2 * (1 - xdep) / (2 + xdep) = 0.958725775
# */
xfd = 0.958725775
xbeta2 = 0.5
# float pl[5];
# double fs01, fs02, fs0, fs1, fs2;
as0 = [
0.33243832, 0.16285370, -0.30924818, -0.10324388, 0.11493334,
-6.777104e-02, 1.577425e-03, -1.240906e-02, 3.241678e-02, -3.503695e-02
]
as1 = [0.19666292, -5.439061e-02]
as2 = [0.14545937, -2.910845e-02]
# float phios, xcos1, xcos2, xcos3;
# float xph1, xph2, xph3, xitm1, xitm2;
# float xlntaur, xitot1, xitot2, xitot3;
# int i, ib;
xph1 = 1.0 + (3.0 * mus * mus - 1.0) * (3.0 * muv * muv - 1.0) * xfd / 8.0
xph2 = -xfd * xbeta2 * 1.5 * mus * muv * da.sqrt(
1.0 - mus * mus) * da.sqrt(1.0 - muv * muv)
xph3 = xfd * xbeta2 * 0.375 * (1.0 - mus * mus) * (1.0 - muv * muv)
# pl[0] = 1.0
# pl[1] = mus + muv
# pl[2] = mus * muv
# pl[3] = mus * mus + muv * muv
# pl[4] = mus * mus * muv * muv
fs01 = as0[0] + (mus + muv) * as0[1] + (mus * muv) * as0[2] + (
mus * mus + muv * muv) * as0[3] + (mus * mus * muv * muv) * as0[4]
fs02 = as0[5] + (mus + muv) * as0[6] + (mus * muv) * as0[7] + (
mus * mus + muv * muv) * as0[8] + (mus * mus * muv * muv) * as0[9]
# for (i = 0; i < 5; i++) {
# fs01 += (double) (pl[i] * as0[i]);
# fs02 += (double) (pl[i] * as0[5 + i]);
# }
# for refl, (ah2o, bh2o, ao3, tau) in zip(reflectance_bands, coefficients):
# ib = find_coefficient_index(center_wl)
# if ib is None:
# raise ValueError("Can't handle band with wavelength '{}'".format(center_wl))
xlntaur = da.log(taur)
fs0 = fs01 + fs02 * xlntaur
fs1 = as1[0] + xlntaur * as1[1]
fs2 = as2[0] + xlntaur * as2[1]
del xlntaur, fs01, fs02
trdown = da.exp(-taur / mus)
trup = da.exp(-taur / muv)
xitm1 = (1.0 - trdown * trup) / 4.0 / (mus + muv)
xitm2 = (1.0 - trdown) * (1.0 - trup)
xitot1 = xph1 * (xitm1 + xitm2 * fs0)
xitot2 = xph2 * (xitm1 + xitm2 * fs1)
xitot3 = xph3 * (xitm1 + xitm2 * fs2)
del xph1, xph2, xph3, xitm1, xitm2, fs0, fs1, fs2
phios = da.deg2rad(phi + 180.0)
xcos1 = 1.0
xcos2 = da.cos(phios)
xcos3 = da.cos(2.0 * phios)
del phios
rhoray = xitot1 * xcos1 + xitot2 * xcos2 * 2.0 + xitot3 * xcos3 * 2.0
return rhoray, trdown, trup
def predict(args):
# get inclusion regions
include_regions = []
exclude_regions = []
if args.within:
from regions import read_ds9
import tempfile
# kludge because regions cries over "FK5", wants lowercase
with tempfile.NamedTemporaryFile(mode="w") as tmpfile, open(
args.within) as regfile:
tmpfile.write(regfile.read().lower())
tmpfile.flush()
include_regions = read_ds9(tmpfile.name)
log.info("read {} inclusion region(s) from {}".format(
len(include_regions), args.within))
# Import source data from WSClean component list
# See https://sourceforge.net/p/wsclean/wiki/ComponentList
(comp_type, radec, stokes, spec_coeff, ref_freq, log_spec_ind,
gaussian_shape) = import_from_wsclean(args.sky_model,
include_regions=include_regions,
exclude_regions=exclude_regions,
point_only=args.points_only,
num=args.num_sources or None)
# Get the support tables
tables = support_tables(
args, ["FIELD", "DATA_DESCRIPTION", "SPECTRAL_WINDOW", "POLARIZATION"])
field_ds = tables["FIELD"]
ddid_ds = tables["DATA_DESCRIPTION"]
spw_ds = tables["SPECTRAL_WINDOW"]
pol_ds = tables["POLARIZATION"]
frequencies = np.sort(
[spw_ds[dd].CHAN_FREQ.data.flatten() for dd in range(len(spw_ds))])
# cluster sources and refit. This only works for delta scale sources
def __cluster(comp_type, radec, stokes, spec_coeff, ref_freq, log_spec_ind,
gaussian_shape, frequencies):
uniq_radec = np.unique(radec)
ncomp_type = []
nradec = []
nstokes = []
nspec_coef = []
nref_freq = []
nlog_spec_ind = []
ngaussian_shape = []
for urd in uniq_radec:
print comp_type.shape
print radec.shape
deltasel = comp_type[radec == urd] == "POINT"
polyspecsel = np.logical_not(spec_coef[radec == urd])
sel = deltasel & polyspecsel
Is = stokes[sel, 0, None] * frequency[None, :]**0
for jj in range(spec_coeff.shape[1]):
Is += spec_coeff[sel, jj, None] * (
frequency[None, :] / ref_freq[sel, None] - 1)**(jj + 1)
Is = np.sum(
Is, axis=0) # collapse over all the sources at this position
logpolyspecsel = np.logical_not(log_spec_coef[radec == urd])
sel = deltasel & logpolyspecsel
Is = np.log(stokes[sel, 0, None] * frequency[None, :]**0)
for jj in range(spec_coeff.shape[1]):
Is += spec_coeff[sel, jj, None] * da.log(
(frequency[None, :] / ref_freq[sel, None])**(jj + 1))
Is = np.exp(Is)
Islogpoly = np.sum(
Is, axis=0) # collapse over all the sources at this position
popt, pfitvar = curve_fit(
lambda i, a, b, c, d: i + a *
(frequency / ref_freq[0, None] - 1) + b *
(frequency / ref_freq[0, None] - 1)**2 + c *
(frequency / ref_freq[sel, None] - 1)**3 + d *
(frequency / ref_freq[0, None] - 1)**3, frequency,
Ispoly + Islogpoly)
if not np.all(np.isfinite(pfitvar)):
popt[0] = np.sum(stokes[sel, 0, None], axis=0)
popt[1:] = np.inf
log.warn(
"Refitting at position {0:s} failed. Assuming flat spectrum source of {1:.2f} Jy"
.format(radec, popt[0]))
else:
pcov = np.sqrt(np.diag(pfitvar))
log.info(
"New fitted flux {0:.3f} Jy at position {1:s} with covariance {2:s}"
.format(popt[0], radec,
", ".join([str(poptp) for poptp in popt])))
ncomp_type.append("POINT")
nradec.append(urd)
nstokes.append(popt[0])
nspec_coef.append(popt[1:])
nref_freq.append(ref_freq[0])
nlog_spec_ind = 0.0
# add back all the gaussians
sel = comp_type[radec] == "GAUSSIAN"
for rd, stks, spec, ref, lspec, gs in zip(radec[sel], stokes[sel],
spec_coef[sel],
ref_freq[sel],
log_spec_ind[sel],
gaussian_shape[sel]):
ncomp_type.append("GAUSSIAN")
nradec.append(rd)
nstokes.append(stks)
nspec_coef.append(spec)
nref_freq.append(ref)
nlog_spec_ind.append(lspec)
ngaussian_shape.append(gs)
log.info(
"Reduced {0:d} components to {1:d} components through by refitting"
.format(len(comp_type), len(ncomp_type)))
return (np.array(ncomp_type), np.array(nradec), np.array(nstokes),
np.array(nspec_coeff), np.array(nref_freq),
np.array(nlog_spec_ind), np.array(ngaussian_shape))
if not args.dontcluster:
(comp_type, radec, stokes, spec_coeff, ref_freq, log_spec_ind,
gaussian_shape) = __cluster(comp_type, radec, stokes, spec_coeff,
ref_freq, log_spec_ind, gaussian_shape,
frequencies)
# Add output column if it isn't present
ms_rows, ms_datatype = ms_preprocess(args)
# sort out resources
args.row_chunks, args.model_chunks = get_budget(
comp_type.shape[0], ms_rows, max([ss.NUM_CHAN.data for ss in spw_ds]),
max([ss.NUM_CORR.data for ss in pol_ds]), ms_datatype, args)
radec = da.from_array(radec, chunks=(args.model_chunks, 2))
stokes = da.from_array(stokes, chunks=(args.model_chunks, 4))
if np.count_nonzero(comp_type == 'GAUSSIAN') > 0:
gaussian_components = True
gshape_chunks = (args.model_chunks, 3)
gaussian_shape = da.from_array(gaussian_shape, chunks=gshape_chunks)
else:
gaussian_components = False
if args.spectra:
spec_chunks = (args.model_chunks, spec_coeff.shape[1])
spec_coeff = da.from_array(spec_coeff, chunks=spec_chunks)
ref_freq = da.from_array(ref_freq, chunks=(args.model_chunks, ))
# List of write operations
writes = []
# Construct a graph for each DATA_DESC_ID
for xds in xds_from_ms(args.ms,
columns=["UVW", "ANTENNA1", "ANTENNA2", "TIME"],
group_cols=["FIELD_ID", "DATA_DESC_ID"],
chunks={"row": args.row_chunks}):
if xds.attrs['FIELD_ID'] != args.fieldid:
continue
# Extract frequencies from the spectral window associated
# with this data descriptor id
field = field_ds[xds.attrs['FIELD_ID']]
ddid = ddid_ds[xds.attrs['DATA_DESC_ID']]
spw = spw_ds[ddid.SPECTRAL_WINDOW_ID.values]
pol = pol_ds[ddid.POLARIZATION_ID.values]
frequency = spw.CHAN_FREQ.data
corrs = pol.NUM_CORR.values
lm = radec_to_lm(radec, field.PHASE_DIR.data)
if args.exp_sign_convention == 'casa':
uvw = -xds.UVW.data
elif args.exp_sign_convention == 'thompson':
uvw = xds.UVW.data
else:
raise ValueError("Invalid sign convention '%s'" % args.sign)
if args.spectra:
# flux density at reference frequency ...
# ... for logarithmic polynomial functions
if log_spec_ind:
Is = da.log(stokes[:, 0, None]) * frequency[None, :]**0
# ... or for ordinary polynomial functions
else:
Is = stokes[:, 0, None] * frequency[None, :]**0
# additional terms of SED ...
for jj in range(spec_coeff.shape[1]):
# ... for logarithmic polynomial functions
if log_spec_ind:
Is += spec_coeff[:, jj, None] * da.log(
(frequency[None, :] / ref_freq[:, None])**(jj + 1))
# ... or for ordinary polynomial functions
else:
Is += spec_coeff[:, jj, None] * (
frequency[None, :] / ref_freq[:, None] - 1)**(jj + 1)
if log_spec_ind: Is = da.exp(Is)
Qs = da.zeros_like(Is)
Us = da.zeros_like(Is)
Vs = da.zeros_like(Is)
spectrum = da.stack(
[Is, Qs, Us, Vs], axis=-1
) # stack along new axis and make it the last axis of the new array
spectrum = spectrum.rechunk(spectrum.chunks[:2] +
(spectrum.shape[2], ))
log.info('-------------------------------------------')
log.info('Nr sources = {0:d}'.format(stokes.shape[0]))
log.info('-------------------------------------------')
log.info('stokes.shape = {0:}'.format(stokes.shape))
log.info('frequency.shape = {0:}'.format(frequency.shape))
if args.spectra: log.info('Is.shape = {0:}'.format(Is.shape))
if args.spectra:
log.info('spectrum.shape = {0:}'.format(spectrum.shape))
# (source, row, frequency)
phase = phase_delay(lm, uvw, frequency)
# If at least one Gaussian component is present in the component list then all
# sources are modelled as Gaussian components (Delta components have zero width)
if gaussian_components:
phase *= gaussian(uvw, frequency, gaussian_shape)
# (source, frequency, corr_products)
brightness = convert(spectrum if args.spectra else stokes,
["I", "Q", "U", "V"], corr_schema(pol))
log.info('brightness.shape = {0:}'.format(brightness.shape))
log.info('phase.shape = {0:}'.format(phase.shape))
log.info('-------------------------------------------')
log.info('Attempting phase-brightness einsum with "{0:s}"'.format(
einsum_schema(pol, args.spectra)))
# (source, row, frequency, corr_products)
jones = da.einsum(einsum_schema(pol, args.spectra), phase, brightness)
log.info('jones.shape = {0:}'.format(jones.shape))
log.info('-------------------------------------------')
if gaussian_components: log.info('Some Gaussian sources found')
else: log.info('All sources are Delta functions')
log.info('-------------------------------------------')
# Identify time indices
_, time_index = da.unique(xds.TIME.data, return_inverse=True)
# Predict visibilities
vis = predict_vis(time_index, xds.ANTENNA1.data, xds.ANTENNA2.data,
None, jones, None, None, None, None)
# Reshape (2, 2) correlation to shape (4,)
if corrs == 4:
vis = vis.reshape(vis.shape[:2] + (4, ))
# Assign visibilities to MODEL_DATA array on the dataset
model_data = xr.DataArray(vis, dims=["row", "chan", "corr"])
xds = xds.assign(**{args.output_column: model_data})
# Create a write to the table
write = xds_to_table(xds, args.ms, [args.output_column])
# Add to the list of writes
writes.append(write)
# Submit all graph computations in parallel
if args.num_workers:
with ProgressBar(), dask.config.set(num_workers=args.num_workers):
dask.compute(writes)
else:
with ProgressBar():
dask.compute(writes)
def _distance(Z, Y, epsilon):
""" Distance function """
Y = Y + epsilon
return Y.sum(axis=(1, 2)) - da.einsum('ijk,ljk->il', Z, da.log(Y))
def __cluster(comp_type, radec, stokes, spec_coeff, ref_freq, log_spec_ind,
gaussian_shape, frequencies):
uniq_radec = np.unique(radec)
ncomp_type = []
nradec = []
nstokes = []
nspec_coef = []
nref_freq = []
nlog_spec_ind = []
ngaussian_shape = []
for urd in uniq_radec:
print comp_type.shape
print radec.shape
deltasel = comp_type[radec == urd] == "POINT"
polyspecsel = np.logical_not(spec_coef[radec == urd])
sel = deltasel & polyspecsel
Is = stokes[sel, 0, None] * frequency[None, :]**0
for jj in range(spec_coeff.shape[1]):
Is += spec_coeff[sel, jj, None] * (
frequency[None, :] / ref_freq[sel, None] - 1)**(jj + 1)
Is = np.sum(
Is, axis=0) # collapse over all the sources at this position
logpolyspecsel = np.logical_not(log_spec_coef[radec == urd])
sel = deltasel & logpolyspecsel
Is = np.log(stokes[sel, 0, None] * frequency[None, :]**0)
for jj in range(spec_coeff.shape[1]):
Is += spec_coeff[sel, jj, None] * da.log(
(frequency[None, :] / ref_freq[sel, None])**(jj + 1))
Is = np.exp(Is)
Islogpoly = np.sum(
Is, axis=0) # collapse over all the sources at this position
popt, pfitvar = curve_fit(
lambda i, a, b, c, d: i + a *
(frequency / ref_freq[0, None] - 1) + b *
(frequency / ref_freq[0, None] - 1)**2 + c *
(frequency / ref_freq[sel, None] - 1)**3 + d *
(frequency / ref_freq[0, None] - 1)**3, frequency,
Ispoly + Islogpoly)
if not np.all(np.isfinite(pfitvar)):
popt[0] = np.sum(stokes[sel, 0, None], axis=0)
popt[1:] = np.inf
log.warn(
"Refitting at position {0:s} failed. Assuming flat spectrum source of {1:.2f} Jy"
.format(radec, popt[0]))
else:
pcov = np.sqrt(np.diag(pfitvar))
log.info(
"New fitted flux {0:.3f} Jy at position {1:s} with covariance {2:s}"
.format(popt[0], radec,
", ".join([str(poptp) for poptp in popt])))
ncomp_type.append("POINT")
nradec.append(urd)
nstokes.append(popt[0])
nspec_coef.append(popt[1:])
nref_freq.append(ref_freq[0])
nlog_spec_ind = 0.0
# add back all the gaussians
sel = comp_type[radec] == "GAUSSIAN"
for rd, stks, spec, ref, lspec, gs in zip(radec[sel], stokes[sel],
spec_coef[sel],
ref_freq[sel],
log_spec_ind[sel],
gaussian_shape[sel]):
ncomp_type.append("GAUSSIAN")
nradec.append(rd)
nstokes.append(stks)
nspec_coef.append(spec)
nref_freq.append(ref)
nlog_spec_ind.append(lspec)
ngaussian_shape.append(gs)
log.info(
"Reduced {0:d} components to {1:d} components through by refitting"
.format(len(comp_type), len(ncomp_type)))
return (np.array(ncomp_type), np.array(nradec), np.array(nstokes),
np.array(nspec_coeff), np.array(nref_freq),
np.array(nlog_spec_ind), np.array(ngaussian_shape))
