diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 97a81b7c..958faa67 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -17,6 +17,7 @@ jobs: # in __init__. below we'll reinstall for the tests. run: | pip install astropy + pip install frank pip install . - name: Cache/Restore the .mpol folder cache uses: actions/cache@v3 diff --git a/mypy.ini b/mypy.ini index a3007204..19b0da10 100644 --- a/mypy.ini +++ b/mypy.ini @@ -14,5 +14,8 @@ ignore_missing_imports = True [mypy-torchkbnufft.*] ignore_missing_imports = True +[mypy-frank.*] +ignore_missing_imports = True + [mypy-fast_histogram.*] ignore_missing_imports = True \ No newline at end of file diff --git a/setup.py b/setup.py index 6d1a62c4..d8c07d67 100644 --- a/setup.py +++ b/setup.py @@ -33,6 +33,7 @@ def get_version(rel_path): "astropy", "tensorboard", "mypy", + "frank>=1.2.1", ], "docs": [ "sphinx>=2.3.0", @@ -55,10 +56,13 @@ def get_version(rel_path): "pyro-ppl", "arviz[all]" ], + "analysis": [ + "frank>=1.2.1", + ], } EXTRA_REQUIRES["dev"] = ( - EXTRA_REQUIRES["test"] + EXTRA_REQUIRES["docs"] + ["pylint", "black", "pre-commit"] + EXTRA_REQUIRES["test"] + EXTRA_REQUIRES["docs"] + EXTRA_REQUIRES["analysis"] + ["pylint", "black", "pre-commit"] ) diff --git a/src/mpol/geometry.py b/src/mpol/geometry.py index 1ce4e8de..b7438971 100644 --- a/src/mpol/geometry.py +++ b/src/mpol/geometry.py @@ -28,7 +28,6 @@ def flat_to_observer(x, y, omega=None, incl=None, Omega=None): Returns: Two tensors representing ``(X, Y)`` in the observer frame. """ - # Rotation matrices result in a *clockwise* rotation of the axes, as defined using the righthand rule. # For example, looking down the z-axis, a positive angle will rotate the x,y axes clockwise. # A vector in the coordinate system will appear as though it has been rotated counter-clockwise. @@ -70,7 +69,7 @@ def flat_to_observer(x, y, omega=None, incl=None, Omega=None): def observer_to_flat(X, Y, omega=None, incl=None, Omega=None): - """Rotate the from to convert a point in the observer frame (X,Y,Z) to the flat (x,y,z) frame. + """Rotate the frame to convert a point in the observer frame (X,Y,Z) to the flat (x,y,z) frame. It is assumed that the +Z axis points *towards* the observer. The rotation operations are the inverse of the :func:`~mpol.geometry.flat_to_observer` operations. @@ -90,9 +89,8 @@ def observer_to_flat(X, Y, omega=None, incl=None, Omega=None): Omega (torch float tensor): A tensor representing the position angle of the ascending node in [radians]. Default 0.0 Returns: - Two tensors representing ``(x, y)`` in the observer frame. + Two tensors representing ``(x, y)`` in the flat frame. """ - # Rotation matrices result in a *clockwise* rotation of the axes, as defined using the righthand rule. # For example, looking down the z-axis, a positive angle will rotate the x,y axes clockwise. # A vector in the coordinate system will appear as though it has been rotated counter-clockwise. diff --git a/src/mpol/images.py b/src/mpol/images.py index 0185766e..5090fd6a 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -315,4 +315,4 @@ def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): hdul = fits.HDUList([hdu]) hdul.writeto(fname, overwrite=overwrite) - hdul.close() + hdul.close() \ No newline at end of file diff --git a/src/mpol/onedim.py b/src/mpol/onedim.py new file mode 100644 index 00000000..58e83131 --- /dev/null +++ b/src/mpol/onedim.py @@ -0,0 +1,170 @@ +import numpy as np +from mpol.utils import torch2npy + +def radialI(icube, geom, chan=0, bins=None): + r""" + Obtain a 1D (radial) brightness profile I(r) from an image cube. + + Parameters + ---------- + icube : `mpol.images.ImageCube` object + Instance of the MPoL `images.ImageCube` class + geom : dict + Dictionary of source geometry. Keys: + "incl" : float, unit=[deg] + Inclination + "Omega" : float, unit=[deg] + Position angle of the ascending node + "omega" : float, unit=[deg] + Argument of periastron + "dRA" : float, unit=[arcsec] + Phase center offset in right ascension. Positive is west of north. + "dDec" : float, unit=[arcsec] + Phase center offset in declination. + chan : int, default=0 + Channel of the image cube corresponding to the desired image + bins : array, default=None, unit=[arcsec] + Radial bin edges to use in calculating I(r). If None, bins will span + the full image, with widths equal to the hypotenuse of the pixels + + Returns + ------- + bin_centers : array, unit=[arcsec] + Radial coordinates of image at center of `bins` + Is : array, unit=[Jy / arcsec^2] (if `image` has these units) + Azimuthally averaged pixel brightness at `rs` + """ + + # projected Cartesian pixel coordinates [arcsec] + xx, yy = icube.coords.sky_x_centers_2D, icube.coords.sky_y_centers_2D + + # shift image center to source center + xc, yc = xx - geom["dRA"], yy - geom["dDec"] + + # deproject image + cos_PA = np.cos(geom["Omega"] * np.pi / 180) + sin_PA = np.sin(geom["Omega"] * np.pi / 180) + xd = xc * cos_PA - yc * sin_PA + yd = xc * sin_PA + yc * cos_PA + xd /= np.cos(geom["incl"] * np.pi / 180) + + # deprojected radial coordinates + rr = np.ravel(np.hypot(xd, yd)) + + if bins is None: + # choose sensible bin size and range + step = np.hypot(icube.coords.cell_size, icube.coords.cell_size) + bins = np.arange(0.0, np.max((abs(xc.ravel()), abs(yc.ravel()))), step) + + bin_counts, bin_edges = np.histogram(a=rr, bins=bins, weights=None) + + # cumulative binned brightness in each annulus + Is, _ = np.histogram(a=rr, bins=bins, + weights=torch2npy(icube.sky_cube[chan]).ravel() + ) + + # mask empty bins + mask = (bin_counts == 0) + Is = np.ma.masked_where(mask, Is) + + # average binned brightness in each annulus + Is /= bin_counts + + bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2 + + return bin_centers, Is + + +def radialV(fcube, geom, rescale_flux, chan=0, bins=None): + r""" + Obtain the 1D (radial) visibility model V(q) corresponding to a 2D MPoL + image. + + Parameters + ---------- + fcube : `~mpol.fourier.FourierCube` object + Instance of the MPoL `fourier.FourierCube` class + geom : dict + Dictionary of source geometry. Keys: + "incl" : float, unit=[deg] + Inclination + "Omega" : float, unit=[deg] + Position angle of the ascending node + "omega" : float, unit=[deg] + Argument of periastron + "dRA" : float, unit=[arcsec] + Phase center offset in right ascension. Positive is west of north. + "dDec" : float, unit=[arcsec] + Phase center offset in declination + rescale_flux : bool + If True, the visibility amplitudes and weights are rescaled to account + for the difference between the inclined (observed) brightness and the + assumed face-on brightness, assuming the emission is optically thick. + The source's integrated (2D) flux is assumed to be: + :math:`F = \cos(i) \int_r^{r=R}{I(r) 2 \pi r dr}`. + No rescaling would be appropriate in the optically thin limit. + chan : int, default=0 + Channel of the image cube corresponding to the desired image + bins : array, default=None, unit=[k\lambda] + Baseline bin edges to use in calculating V(q). If None, bins will span + the model baseline distribution, with widths equal to the hypotenuse of + the (u, v) coordinates + + Returns + ------- + bin_centers : array, unit=:math:[`k\lambda`] + Baselines corresponding to `u` and `v` + Vs : array, unit=[Jy] + Visibility amplitudes at `q` + + Notes + ----- + This routine requires the `frank `_ package + """ + from frank.geometry import apply_phase_shift, deproject + + # projected model (u,v) points [k\lambda] + uu, vv = fcube.coords.sky_u_centers_2D, fcube.coords.sky_v_centers_2D + + # visibilities + V = torch2npy(fcube.ground_cube[chan]).ravel() + + # phase-shift the visibilities + Vp = apply_phase_shift(uu.ravel() * 1e3, vv.ravel() * 1e3, V, geom["dRA"], + geom["dDec"], inverse=True) + + # deproject the (u,v) points + up, vp, _ = deproject(uu.ravel() * 1e3, vv.ravel() * 1e3, geom["incl"], + geom["Omega"]) + + # if the source is optically thick, rescale the deprojected V(q) + if rescale_flux: + Vp.real /= np.cos(geom["incl"] * np.pi / 180) + + # convert back to [k\lambda] + up /= 1e3 + vp /= 1e3 + + # deprojected baselines + qq = np.hypot(up, vp) + + if bins is None: + # choose sensible bin size and range + step = np.hypot(fcube.coords.du, fcube.coords.dv) / 2 + bins = np.arange(0.0, max(qq), step) + + bin_counts, bin_edges = np.histogram(a=qq, bins=bins, weights=None) + + # cumulative binned visibility amplitude in each annulus + Vs, _ = np.histogram(a=qq, bins=bins, weights=Vp) + + # mask empty bins + mask = (bin_counts == 0) + Vs = np.ma.masked_where(mask, Vs) + + # average binned visibility amplitude in each annulus + Vs /= bin_counts + + bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2 + + return bin_centers, Vs \ No newline at end of file diff --git a/src/mpol/utils.py b/src/mpol/utils.py index ac8ced23..92096b19 100644 --- a/src/mpol/utils.py +++ b/src/mpol/utils.py @@ -302,7 +302,6 @@ def get_optimal_image_properties(image_width, u, v): npix : int Number of pixels of cell_size to equal (or slightly exceed) the image width (npix will be rounded up and enforced as even). - Notes ----- No assumption or correction is made concerning whether the spatial diff --git a/test/conftest.py b/test/conftest.py index 86e8827a..39eee5e6 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,8 +1,9 @@ import numpy as np import pytest +import torch from astropy.utils.data import download_file -from mpol import coordinates, gridding +from mpol import coordinates, fourier, gridding, images from mpol.__init__ import zenodo_record # We need a fixture which provides mock visibilities of the sort we'd @@ -116,6 +117,71 @@ def dataset_cont(mock_visibility_data_cont, coords): return averager.to_pytorch_dataset() +@pytest.fixture(scope="session") +def mock_1d_archive(): + # use astropy routines to cache data + fname = download_file( + f"https://zenodo.org/record/{zenodo_record}/files/mock_disk_1d.npz", + cache=True, + pkgname="mpol", + ) + + return np.load(fname, allow_pickle=True) + + +@pytest.fixture +def mock_1d_image_model(mock_1d_archive): + m = mock_1d_archive + rtrue = m['rtrue'] + itrue = m['itrue'] + i2dtrue = m['i2dtrue'] + xmax = ymax = m['xmax'] + geom = m['geometry'] + geom = geom[()] + + coords = coordinates.GridCoords(cell_size=xmax * 2 / i2dtrue.shape[0], + npix=i2dtrue.shape[0]) + + # the center of the array is already at the center of the image --> + # undo this as expected by input to ImageCube + i2dtrue = np.flip(np.fft.fftshift(i2dtrue), 1) + + # pack the numpy image array into an ImageCube + packed_cube = np.broadcast_to(i2dtrue, (1, coords.npix, coords.npix)).copy() + packed_tensor = torch.from_numpy(packed_cube) + cube_true = images.ImageCube(coords=coords, nchan=1, cube=packed_tensor) + + return rtrue, itrue, cube_true, xmax, ymax, geom + + +@pytest.fixture +def mock_1d_vis_model(mock_1d_archive): + m = mock_1d_archive + Vtrue = m['vis'] + Vtrue_dep = m['vis_dep'] + q_dep = m['baselines_dep'] + geom = m['geometry'] + geom = geom[()] + + xmax = m['xmax'] + i2dtrue = m['i2dtrue'] + + coords = coordinates.GridCoords(cell_size=xmax * 2 / i2dtrue.shape[0], + npix=i2dtrue.shape[0]) + + # create a FourierCube + packed_cube = np.broadcast_to(Vtrue, (1, len(Vtrue))).copy() + packed_tensor = torch.from_numpy(packed_cube) + cube_true = fourier.FourierCube(coords=coords) + + # insert the vis tensor into the FourierCube ('vis' would typically be + # populated by taking the FFT of an image) + cube_true.ground_cube = packed_tensor + + return cube_true, Vtrue_dep, q_dep, geom + + + @pytest.fixture def crossvalidation_products(mock_visibility_data): # test the crossvalidation with a smaller set of image / Fourier coordinates than normal, diff --git a/test/onedim_test.py b/test/onedim_test.py new file mode 100644 index 00000000..b78cc5ba --- /dev/null +++ b/test/onedim_test.py @@ -0,0 +1,118 @@ +import matplotlib.pyplot as plt +import numpy as np + +from mpol.onedim import radialI, radialV +from mpol.plot import plot_image +from mpol.utils import torch2npy + +def test_radialI(mock_1d_image_model, tmp_path): + # obtain a 1d radial brightness profile I(r) from an image + + rtrue, itrue, icube, _, _, geom = mock_1d_image_model + + bins = np.linspace(0, 2.0, 100) + + rtest, itest = radialI(icube, geom, bins=bins) + + fig, ax = plt.subplots(ncols=2, figsize=(10,5)) + + plot_image(np.squeeze(torch2npy(icube.sky_cube)), extent=icube.coords.img_ext, + ax=ax[0], clab='Jy / sr') + + ax[1].plot(rtrue, itrue, 'k', label='truth') + ax[1].plot(rtest, itest, 'r.-', label='recovery') + + ax[0].set_title(f"Geometry:\n{geom}", fontsize=7) + + ax[1].set_xlabel('r [arcsec]') + ax[1].set_ylabel('I [Jy / sr]') + ax[1].legend() + + fig.savefig(tmp_path / "test_radialI.png", dpi=300) + plt.close("all") + + expected = [ + 6.40747314e+10, 4.01920507e+10, 1.44803534e+10, 2.94238627e+09, + 1.28782935e+10, 2.68613199e+10, 2.26564596e+10, 1.81151845e+10, + 1.52128965e+10, 1.05640352e+10, 1.33411204e+10, 1.61124502e+10, + 1.41500539e+10, 1.20121195e+10, 1.11770326e+10, 1.19676913e+10, + 1.20941686e+10, 1.09498286e+10, 9.74236410e+09, 7.99589196e+09, + 5.94787809e+09, 3.82074946e+09, 1.80823933e+09, 4.48414819e+08, + 3.17808840e+08, 5.77317876e+08, 3.98851281e+08, 8.06459834e+08, + 2.88706161e+09, 6.09577814e+09, 6.98556762e+09, 4.47436415e+09, + 1.89511273e+09, 5.96604356e+08, 3.44571640e+08, 5.65906765e+08, + 2.85854589e+08, 2.67589013e+08, 3.98357054e+08, 2.97052261e+08, + 3.82744591e+08, 3.52239791e+08, 2.74336969e+08, 2.28425747e+08, + 1.82290043e+08, 3.16077299e+08, 1.18465538e+09, 3.32239287e+09, + 5.26718846e+09, 5.16458748e+09, 3.58114198e+09, 2.13431954e+09, + 1.40936556e+09, 1.04032244e+09, 9.24050422e+08, 8.46829316e+08, + 6.80909295e+08, 6.83812465e+08, 6.91856237e+08, 5.29227136e+08, + 3.97557293e+08, 3.54893419e+08, 2.60997039e+08, 2.09306498e+08, + 1.93930693e+08, 6.97032407e+07, 6.66090083e+07, 1.40079594e+08, + 7.21775931e+07, 3.23902663e+07, 3.35932300e+07, 7.63318789e+06, + 1.29740981e+07, 1.44300351e+07, 8.06249624e+06, 5.85567843e+06, + 1.42637174e+06, 3.21445075e+06, 1.83763663e+06, 1.16926652e+07, + 2.46918188e+07, 1.60206523e+07, 3.26596592e+06, 1.27837319e+05, + 2.27104612e+04, 4.77267063e+03, 2.90467640e+03, 2.88482230e+03, + 1.43402521e+03, 1.54791996e+03, 7.23397046e+02, 1.02561351e+03, + 5.24845888e+02, 1.47320552e+03, 7.40419174e+02, 4.59029378e-03, + 0.00000000e+00, 0.00000000e+00, 0.00000000e+00 + ] + + np.testing.assert_allclose(itest, expected, rtol=1e-6, + err_msg="test_radialI") + + +def test_radialV(mock_1d_vis_model, tmp_path): + # obtain a 1d radial visibility profile V(q) from 2d visibilities + + fcube, Vtrue_dep, q_dep, geom = mock_1d_vis_model + + bins = np.linspace(1,5e3,100) + + qtest, Vtest = radialV(fcube, geom, rescale_flux=True, bins=bins) + + fig, ax = plt.subplots() + + ax.plot(q_dep / 1e6, Vtrue_dep, 'k.', label='truth deprojected') + ax.plot(qtest / 1e3, Vtest, 'r.-', label='recovery') + + ax.set_xlim(-0.5, 6) + ax.set_xlabel(r'Baseline [M$\lambda$]') + ax.set_ylabel('Re(V) [Jy]') + ax.set_title(f"Geometry {geom}", fontsize=10) + ax.legend() + + fig.savefig(tmp_path / "test_radialV.png", dpi=300) + plt.close("all") + + expected = [ + 2.53998336e-01, 1.59897580e-01, 8.59460326e-02, 7.42189236e-02, + 5.75440687e-02, 1.81324892e-02, -2.92922689e-03, 3.14354163e-03, + 6.72339399e-03, -8.54632390e-03, -1.73385166e-02, -4.03826092e-03, + 1.45595908e-02, 1.61681713e-02, 5.93475866e-03, 2.45555912e-04, + -7.05014619e-04, -6.09266430e-03, -1.02454088e-02, -2.80944776e-03, + 8.58212558e-03, 8.39132158e-03, -6.52523293e-04, -4.34778158e-03, + 1.08035980e-04, 3.40903070e-03, 2.26682041e-03, 2.42465437e-03, + 5.07968926e-03, 4.83377443e-03, 1.26300648e-03, 1.11930639e-03, + 6.45157513e-03, 1.05751150e-02, 9.14016956e-03, 5.92209210e-03, + 5.18455986e-03, 5.88802559e-03, 5.49315770e-03, 4.96398638e-03, + 5.81115311e-03, 5.95304063e-03, 3.16208083e-03, -1.71765038e-04, + -4.64532628e-04, 1.12448670e-03, 1.84297313e-03, 1.48094594e-03, + 1.12953770e-03, 1.01370816e-03, 6.57047907e-04, 1.37570722e-04, + 3.00648884e-04, 1.04847404e-03, 1.16856102e-03, 3.08940761e-04, + -5.65721897e-04, -8.38907531e-04, -8.71976125e-04, -1.09567680e-03, + -1.42077854e-03, -1.33702627e-03, -9.96839047e-04, -1.16400192e-03, + -1.43584618e-03, -1.07454472e-03, -6.44900590e-04, -4.86165342e-04, + -1.96851463e-04, 5.04190986e-05, 5.73950179e-05, 2.79905736e-04, + 7.52685847e-04, 1.12546048e-03, 1.37149548e-03, 1.35835560e-03, + 1.06470794e-03, 8.81423014e-04, 8.76827161e-04, 9.03579902e-04, + 8.39818807e-04, 5.19936424e-04, 1.46415537e-04, 3.29054769e-05, + 7.30096312e-05, 6.47553400e-05, 2.18817382e-05, 4.47955432e-06, + 7.34705616e-06, 9.06184045e-06, 9.45269846e-06, 1.00464939e-05, + 8.28166011e-06, 7.09361681e-06, 6.43221021e-06, 3.12425880e-06, + 2.57495214e-07, 6.48560373e-07, 1.88421498e-07 + ] + + np.testing.assert_allclose(Vtest, expected, rtol=1e-6, + err_msg="test_radialV") \ No newline at end of file