def test_convolve2gaussres(nx, ny, nband, alpha):
freq = np.linspace(0.5e9, 1.5e9, nband)
ref_freq = freq[0]
Gausspari = ()
es = np.linspace(15, 5, nband)
for v in range(nband):
Gausspari += ((es[v], es[v], 0.0),)
x = np.arange(-nx/2, nx/2)
y = np.arange(-ny/2, ny/2)
xx, yy = np.meshgrid(x, y, indexing='ij')
restored = np.zeros((nband, nx, ny))
for v in range(nband):
restored[v] = Gaussian2D(xx, yy, Gausspari[v], normalise=False) * (freq[v]/ref_freq)**alpha
conv_model, _ = convolve2gaussres(restored, xx, yy, Gausspari[0], 8, gausspari=Gausspari)
I = np.argwhere(conv_model[-1] > 0.05).squeeze()
Ix = I[:, 0]
Iy = I[:, 1]
comps = conv_model[:, Ix, Iy]
weights = np.ones((nband))
out = fit_spi_components(comps.T, weights, freq, ref_freq)
assert_allclose(1+alpha, 1+out[0, :]) # offset for relative difference
assert_allclose(out[2, :], restored[0, Ix, Iy])
def test_dask_fit_spi_components_vs_np():
from africanus.model.spi import fit_spi_components as np_fit_spi
from africanus.model.spi.dask import fit_spi_components
da = pytest.importorskip("dask.array")
np.random.seed(123)
ncomps = 800
alphas = -0.7 + 0.25 * np.random.randn(ncomps, 1)
i0s = 5.0 + np.random.randn(ncomps, 1)
nfreqs = 1000
freqs = np.linspace(0.5, 1.5, nfreqs).reshape(1, nfreqs)
freq0 = 0.7
model = i0s * (freqs / freq0)**alphas
sigma = np.abs(0.25 + 0.1 * np.random.randn(nfreqs))
data = model + sigma[None, :] * np.random.randn(ncomps, nfreqs)
weights = 1.0 / sigma**2
freqs = freqs.squeeze()
alpha1, alphavar1, I01, I0var1 = np_fit_spi(data, weights, freqs, freq0)
# now for the dask version
data_dask = da.from_array(data, chunks=(100, nfreqs))
weights_dask = da.from_array(weights, chunks=(nfreqs))
freqs_dask = da.from_array(freqs, chunks=(nfreqs))
alpha2, alphavar2, I02, I0var2 = fit_spi_components(
data_dask, weights_dask, freqs_dask, freq0).compute()
np.testing.assert_array_almost_equal(alpha1, alpha2, decimal=10)
np.testing.assert_array_almost_equal(alphavar1, alphavar2, decimal=10)
np.testing.assert_array_almost_equal(I01, I02, decimal=10)
np.testing.assert_array_almost_equal(I0var1, I0var2, decimal=10)
def test_fit_spi_components_vs_scipy():
"""
Here we just test the per component spi fitter against
a looped version of scipy's curve_fit
:return:
"""
from africanus.model.spi import fit_spi_components
curve_fit = pytest.importorskip("scipy.optimize").curve_fit
np.random.seed(123)
ncomps = 250
alphas = -0.7 + 0.25 * np.random.randn(ncomps, 1)
i0s = 5.0 + np.random.randn(ncomps, 1)
nfreqs = 100
freqs = np.linspace(0.5, 1.5, nfreqs).reshape(1, nfreqs)
freq0 = 0.7
beams = np.zeros((ncomps, nfreqs))
for i in range(ncomps):
beams[i, :] = np.sinc(freqs - freqs[0])
model = beams * i0s * (freqs / freq0)**alphas
sigma = np.abs(0.25 + 0.1 * np.random.randn(nfreqs))
data = model + sigma[None, :] * np.random.randn(ncomps, nfreqs)
weights = 1.0 / sigma**2
alpha1, alphavar1, I01, I0var1 = fit_spi_components(data,
weights,
freqs.squeeze(),
freq0,
tol=1e-8)
def spi_func(nu, I0, alpha, beam=1.0):
return beam * I0 * nu**alpha
I02 = np.zeros(ncomps)
I0var2 = np.zeros(ncomps)
alpha2 = np.zeros(ncomps)
alphavar2 = np.zeros(ncomps)
for i in range(ncomps):
def fit_func(nu, I0, alpha):
return spi_func(nu, I0, alpha, beam=beams[i])
popt, pcov = curve_fit(fit_func, (freqs / freq0).squeeze(),
data[i, :],
sigma=np.diag(sigma**2),
p0=np.array([1.0, -0.7]),
absolute_sigma=False)
I02[i] = popt[0]
I0var2[i] = pcov[0, 0]
alpha2[i] = popt[1]
alphavar2[i] = pcov[1, 1]
np.testing.assert_array_almost_equal(alpha1, alpha2, decimal=5)
np.testing.assert_array_almost_equal(alphavar1, alphavar2, decimal=5)
np.testing.assert_array_almost_equal(I01, I02, decimal=5)
np.testing.assert_array_almost_equal(I0var1, I0var2, decimal=5)
def test_fit_spi_components_vs_scipy():
"""
Here we just test the per component spi fitter against
a looped version of scipy's curve_fit
:return:
"""
from africanus.model.spi import fit_spi_components
curve_fit = pytest.importorskip("scipy.optimize").curve_fit
np.random.seed(123)
ncomps = 25
alphas = -0.7 + 0.25 * np.random.randn(ncomps, 1)
i0s = 5.0 + np.random.randn(ncomps, 1)
nfreqs = 100
freqs = np.linspace(0.5, 1.5, nfreqs).reshape(1, nfreqs)
freq0 = 0.7
model = i0s * (freqs / freq0)**alphas
sigma = np.abs(0.25 + 0.1 * np.random.randn(nfreqs))
data = model + sigma[None, :] * np.random.randn(ncomps, nfreqs)
weights = 1.0 / sigma**2
alpha1, alphavar1, I01, I0var1 = fit_spi_components(data,
weights,
freqs.squeeze(),
freq0,
tol=1e-8)
def spi_func(nu, I0, alpha):
return I0 * nu**alpha
I02 = np.zeros(ncomps)
I0var2 = np.zeros(ncomps)
alpha2 = np.zeros(ncomps)
alphavar2 = np.zeros(ncomps)
for i in range(ncomps):
popt, pcov = curve_fit(spi_func, (freqs / freq0).squeeze(),
data[i, :],
sigma=np.diag(sigma**2),
p0=np.array([1.0, -0.7]))
I02[i] = popt[0]
I0var2[i] = pcov[0, 0]
alpha2[i] = popt[1]
alphavar2[i] = pcov[1, 1]
np.testing.assert_array_almost_equal(alpha1, alpha2, decimal=6)
# note variances not necessarily accurate to within tol because
# scipy uses LM instead of GN
np.testing.assert_array_almost_equal(alphavar1, alphavar2, decimal=3)
np.testing.assert_array_almost_equal(I01, I02, decimal=6)
np.testing.assert_array_almost_equal(I0var1, I0var2, decimal=3)