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))