def test_spectral_model(spectral_model_inputs, freq):
I, spi, log_si, ref_freq = spectral_model_inputs
# Compute spectral model with numpy implementations
ordinary_spec_model = ordinary_spectral_model(I, spi, log_si, freq,
ref_freq)
log_spec_model = log_spectral_model(I, spi, log_si, freq, ref_freq)
# Choose between ordinary and log spectral index
# based on log_si array
spec_model = np.where(
log_si[:, None] == True, # noqa
log_spec_model,
ordinary_spec_model)
# Compare with our implementation
model = spectra(I, spi, log_si, ref_freq, freq)
assert_array_almost_equal(model, spec_model)
# Ensure positive flux
posI = I.copy()
posI[I < 0.0] = -posI[I < 0.0]
log_si_all_true = np.full(posI.shape, True, dtype=np.bool)
log_spec_model = log_spectral_model(posI, spi, log_si_all_true, freq,
ref_freq)
model = spectra(posI, spi, True, ref_freq, freq)
assert_array_almost_equal(model, log_spec_model)
model = spectra(I, spi, False, ref_freq, freq)
assert_array_almost_equal(model, ordinary_spec_model)
def test_spectral_model(spectral_model_inputs, freq):
I, spi, log_si, ref_freq = spectral_model_inputs
# Ensure positive flux for logarithmic polynomials
I[log_si] = np.abs(I[log_si])
spi[log_si] = np.abs(spi[log_si])
# Compute spectral model with numpy implementations
ordinary_spec_model = ordinary_spectral_model(I, spi, log_si,
freq, ref_freq)
log_spec_model = log_spectral_model(I, spi, log_si,
freq, ref_freq)
# Choose between ordinary and log spectral index
# based on log_si array
spec_model = np.where(log_si[:, None] == True, # noqa
log_spec_model,
ordinary_spec_model)
# Compare with our implementation
model = spectra(I, spi, log_si, ref_freq, freq)
assert_array_almost_equal(model, spec_model)
model = spectra(I, spi, False, ref_freq, freq)
assert_array_almost_equal(model, ordinary_spec_model)
log_si_all_true = np.full(I.shape, True, dtype=np.bool)
I = np.abs(I) # noqa
spi = np.abs(spi)
log_spec_model = log_spectral_model(I, spi, log_si_all_true,
freq, ref_freq)
model = spectra(I, spi, True, ref_freq, freq)
assert_array_almost_equal(model, log_spec_model)
def test_wsclean_predict(chunks):
row = sum(chunks['rows'])
src = sum(chunks['source'])
chan = sum(chunks['channels'])
rs = np.random.RandomState(42)
source_sel = rs.randint(0, 2, src).astype(np.bool)
source_type = np.where(source_sel, "POINT", "GAUSSIAN")
gauss_shape = rs.normal(size=(src, 3))
uvw = rs.normal(size=(row, 3))
lm = rs.normal(size=(src, 2)) * 1e-5
flux = rs.normal(size=src)
coeffs = rs.normal(size=(src, 2))
log_poly = rs.randint(0, 2, src, dtype=np.bool)
flux[log_poly] = np.abs(flux[log_poly])
coeffs[log_poly] = np.abs(coeffs[log_poly])
freq = np.linspace(.856e9, 2 * .856e9, chan)
ref_freq = np.full(src, freq[freq.shape[0] // 2])
# WSClean visibilities
vis = wsclean_predict(uvw, lm, source_type, flux, coeffs, log_poly,
ref_freq, gauss_shape, freq)
# Compute it another way. Note the CASA coordinate convention
# used by wsclean
phase = phase_delay(lm, uvw, freq, convention='casa')
spectrum = spectra(flux, coeffs, log_poly, ref_freq, freq)
shape = gaussian(uvw, freq, gauss_shape)
# point sources don't' contribute to the shape
shape[source_sel] = 1.0
np_vis = np.einsum("srf,srf,sf->rf", shape, phase, spectrum)[:, :, None]
assert_almost_equal(np_vis, vis)
def impl(uvw, lm, source_type, flux, coeffs, log_poly, ref_freq,
gauss_shape, frequency):
spectrum = spectra(flux, coeffs, log_poly, ref_freq, frequency)
return wsclean_predict_impl(uvw, lm, source_type, gauss_shape,
frequency, spectrum, dtype)