Merge branch 'main' into plot_routines

.pre-commit-config.yaml

@@ -2,7 +2,7 @@
 # See https://pre-commit.com/hooks.html for more hooks
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v3.4.0
+    rev: v4.4.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
@@ -12,11 +12,11 @@ repos:
       - id: detect-private-key
       - id: name-tests-test
   - repo: https://github.com/psf/black
-    rev: 22.3.0
+    rev: 23.1.0
     hooks:
       - id: black
   - repo: https://github.com/PyCQA/isort
-    rev: 5.8.0
+    rev: 5.12.0
     hooks:
       - id: isort
         args: []

docs/api.rst

@@ -62,16 +62,6 @@ Training and testing
 .. automodule:: mpol.training
 
 
-Connectors
-----------
-
-The objects in the Images and Precomposed modules are focused on bringing some image-plane model to the space of the data, where the similarity of the model visibilities to the data visibilities will be evaluated by a negative log-likelihood loss. In some situations, though, it is useful to have access to the residual visibilities directly. For example, for visualization or debugging purposes.
-
-Connectors are a PyTorch layer to help compute those residual visibilities (on a gridded form).
-
-.. automodule:: mpol.connectors
-
-
 Cross-validation
 ----------------
 

docs/changelog.md

@@ -2,6 +2,11 @@
 
 # Changelog
 
+## v0.1.4
+
+- Removed the `GriddedResidualConnector` class and the `src/connectors.py` module. Moved `index_vis` to `datasets.py`.
+- Changed BaseCube, ImageCube, and FourierCube initialization signatures
+
 ## v0.1.3
 
 - Added the {func}`mpol.fourier.make_fake_data` routine and the [Mock Data tutorial](ci-tutorials/fakedata.md).

docs/ci-tutorials/crossvalidation.md

@@ -39,7 +39,6 @@ from astropy.utils.data import download_file
 from torch.utils.tensorboard import SummaryWriter
 
 from mpol import (
-    connectors,
     coordinates,
     crossval,
     datasets,
@@ -72,18 +71,25 @@ data_im = np.imag(data)
 # define the image dimensions, making sure they are big enough to fit all
 # of the expected emission
 coords = coordinates.GridCoords(cell_size=0.03, npix=180)
-gridder = gridding.Gridder(
+averager = gridding.DataAverager(
     coords=coords, uu=uu, vv=vv, weight=weight, data_re=data_re, data_im=data_im
 )
 
 # export to PyTorch dataset
-dset = gridder.to_pytorch_dataset()
+dset = averager.to_pytorch_dataset()
 ```
 
+Now, let's also make a diagnostic dirty image
+
 ```{code-cell}
+imager = gridding.DirtyImager(
+    coords=coords, uu=uu, vv=vv, weight=weight, data_re=data_re, data_im=data_im
+)
+
+
 # Show the dirty image
-img, beam = gridder.get_dirty_image(weighting="briggs", robust=0.0)
-kw = {"origin": "lower", "extent": gridder.coords.img_ext}
+img, beam = imager.get_dirty_image(weighting="briggs", robust=0.0)
+kw = {"origin": "lower", "extent": imager.coords.img_ext}
 fig, ax = plt.subplots(ncols=1)
 ax.imshow(np.squeeze(img), **kw)
 ax.set_title("image")
@@ -160,7 +166,7 @@ As you can see, this training set looks very similar to the full dataset, with t
 
 It turns out that the missing holes in the real dataset are quite important to image fidelity---if we had complete $u$,$v$ coverage, we wouldn't need to be worrying about CLEAN or RML imaging techniques in the first place! When we make a new interferometric observation, it will have it's own (different) set of missing holes depending on array configuration, observation duration, and hour angle coverage. We would like our cross validation slices to simulate the $u$,$v$ distribution of possible *new datasets*, and, at least for ALMA, random sampling doesn't probe this very well.
 
-Instead, we suggest an approach where we break the UV plane into radial ($q=\sqrt{u^2 + v^2}$) and azimuthal ($\phi = \mathrm{arctan2}(v,u)$) cells and cross validate by drawing a $K$-fold subselection of these cells. This is just one potential suggestion. There are, of course, no limits on how you might split your dataset for cross-validation; it really depends on what works best for your imaging goals.
+Instead, we suggest an approach where we break the UV plane into radial ($q=\sqrt{u^2 + v^2}$) and azimuthal ($\phi = \mathrm{arctan2}(v,u)$) cells and cross validate by drawing a $K$-fold sub-selection of these cells. This is just one potential suggestion. There are, of course, no limits on how you might split your dataset for cross-validation; it really depends on what works best for your imaging goals.
 
 ```{code-cell}
 # create a radial and azimuthal partition
@@ -178,49 +184,37 @@ k_fold_datasets = [(train, test) for (train, test) in cv]
 
 ```{code-cell}
 flayer = fourier.FourierCube(coords=coords)
-flayer.forward(torch.zeros(dset.nchan, coords.npix, coords.npix))
+flayer(torch.zeros(dset.nchan, coords.npix, coords.npix))
 ```
 
-The following plots visualize how we've split up the data. For each $K$-fold, we have the "training" visibilities, the dirty image corresponding to those training visibilities, and the "test" visibilities which will be used to evaluate the predictive ability of the model.
+The following plots visualize how we've split up the data. For each $K$-fold, we have the "training" visibilities and the "test" visibilities which will be used to evaluate the predictive ability of the model.
 
 ```{code-cell}
-fig, ax = plt.subplots(nrows=k, ncols=3, figsize=(6, 10))
+fig, ax = plt.subplots(nrows=k, ncols=2, figsize=(4, 10))
 
 for i, (train_subset, test_subset) in enumerate(k_fold_datasets):
 
-    rtrain = connectors.GriddedResidualConnector(flayer, train_subset)
-    rtrain.forward()
-    rtest = connectors.GriddedResidualConnector(flayer, test_subset)
-    rtest.forward()
-
-    vis_ext = rtrain.coords.vis_ext
-    img_ext = rtrain.coords.img_ext
+    # train_subset and test_subset are `GriddedDataset`s
 
-    train_mask = rtrain.ground_mask[0]
-    train_chan = rtrain.sky_cube[0]
-
-    test_mask = rtest.ground_mask[0]
-    test_chan = rtest.sky_cube[0]
+    train_mask = train_subset.ground_mask[0]
+    test_mask = test_subset.ground_mask[0]
 
     ax[i, 0].imshow(
         train_mask.detach().numpy(),
         interpolation="none",
         origin="lower",
-        extent=vis_ext,
+        extent=coords.vis_ext,
         cmap="GnBu",
     )
 
-    ax[i, 1].imshow(train_chan.detach().numpy(), origin="lower", extent=img_ext)
-
-    ax[i, 2].imshow(
-        test_mask.detach().numpy(), origin="lower", extent=vis_ext, cmap="GnBu"
+    ax[i, 1].imshow(
+        test_mask.detach().numpy(), origin="lower", extent=coords.vis_ext, cmap="GnBu"
     )
 
     ax[i, 0].set_ylabel("k-fold {:}".format(i))
 
 ax[0, 0].set_title("train mask")
-ax[0, 1].set_title("train dirty img.")
-ax[0, 2].set_title("test mask")
+ax[0, 1].set_title("test mask")
 
 for a in ax.flatten():
     a.xaxis.set_ticklabels([])
@@ -246,7 +240,7 @@ def train(model, dset, config, optimizer, writer=None):
         model.zero_grad()
 
         # get the predicted model
-        vis = model.forward()
+        vis = model()
 
         # get the sky cube too
         sky_cube = model.icube.sky_cube
@@ -274,7 +268,7 @@ We also create a separate "test" function to evaluate the trained model against
 def test(model, dset):
     model.train(False)
     # evaluate test score
-    vis = model.forward()
+    vis = model()
     loss = losses.nll_gridded(vis, dset)
     return loss.item()
 ```

docs/ci-tutorials/fakedata.md

@@ -259,7 +259,7 @@ image = ImageCube.from_image_properties(cell_size=cell_size, npix=npix, nchan=1,
 If you want to double-check that the image was correctly inserted, you can do
 ```
 # double check it went in correctly
-plt.imshow(np.squeeze(utils.packed_cube_to_sky_cube(image.forward()).detach().numpy()), origin="lower")
+plt.imshow(np.squeeze(utils.packed_cube_to_sky_cube(image()).detach().numpy()), origin="lower")
 ```
 to see that it's upright and not flipped.
 
@@ -335,7 +335,7 @@ from mpol import coordinates, gridding
 # well set the
 coords = coordinates.GridCoords(cell_size=cell_size, npix=npix)
 
-gridder = gridding.Gridder(
+imager = gridding.DirtyImager(
     coords=coords,
     uu=uu,
     vv=vv,
@@ -352,12 +352,12 @@ print(noise_estimate, "Jy / dirty beam")
 ```
 
 ```{code-cell} ipython3
-img, beam = gridder.get_dirty_image(weighting="briggs", robust=1.0, unit="Jy/arcsec^2")
+img, beam = imager.get_dirty_image(weighting="briggs", robust=1.0, unit="Jy/arcsec^2")
 ```
 
 ```{code-cell} ipython3
 chan = 0
-kw = {"origin": "lower", "interpolation": "none", "extent": gridder.coords.img_ext}
+kw = {"origin": "lower", "interpolation": "none", "extent": imager.coords.img_ext}
 fig, ax = plt.subplots(ncols=2, figsize=(6.0, 4))
 ax[0].imshow(beam[chan], **kw)
 ax[0].set_title("beam")
@@ -373,7 +373,7 @@ We can even subtract this on a pixel-by-pixel basis and compare to the original
 
 ```{code-cell} ipython3
 chan = 0
-kw = {"origin": "lower", "interpolation": "none", "extent": gridder.coords.img_ext}
+kw = {"origin": "lower", "interpolation": "none", "extent": imager.coords.img_ext}
 fig, ax = plt.subplots(ncols=3, figsize=(6.0, 3))
 
 ax[0].imshow(flux_scaled[chan], **kw)

docs/ci-tutorials/gpu_setup.rst

@@ -161,7 +161,7 @@ step our optimizer.
         model.zero_grad()
 
         # forward pass
-        vis = model.forward()
+        vis = model()
 
         # get skycube from our forward model
         sky_cube = model.icube.sky_cube

docs/ci-tutorials/gridder.md

@@ -72,7 +72,6 @@ print("Dataset has {:} visibilities".format(nvis))
 
 Therefore, understand that the following baseline and visibility scatter plots are showing about a third of a million points.
 
-
 Here, we'll plot the baselines corresponding to the first channel of the dataset by simply marking a point for every spatial frequency coordinate, $u$ and $v$, in the dataset
 
 ```{code-cell}
@@ -114,7 +113,7 @@ ax[3].set_xlabel(r"$q$ [k$\lambda$]");
 There are nearly a third of a million points in each figure, and each is quite noisy, so we can't learn much from this plot alone. But we should be reassured that we see similar types of scatter as we might observe were we to inspect the raw data using CASA's [plotms](https://casadocs.readthedocs.io/en/v6.5.2/api/tt/casaplotms.plotms.html?highlight=plotms) tool.
 
 
-## The {class}`mpol.coordinates.GridCoords` object
+## The {class}`~mpol.coordinates.GridCoords` object
 
 Now, lets familiarize ourselves with MPoL's {class}`mpol.coordinates.GridCoords` object.
 
@@ -128,7 +127,7 @@ Two numbers, `cell_size` and `npix`, uniquely define a grid in image space and i
 coords = coordinates.GridCoords(cell_size=0.005, npix=800)
 ```
 
-The GridCoords object is mainly a container for all of the information about this grid. You can see all of the properties accessible in the {py:class}`mpol.coordinates.GridCoords` API documentation. The information you'll most likely want to access are the image dimensions
+The {class}`mpol.coordinates.GridCoords` object is mainly a container for all of the information about this grid. You can see all of the properties accessible in the {py:class}`mpol.coordinates.GridCoords` API documentation. The information you'll most likely want to access are the image dimensions
 
 ```{code-cell}
 coords.img_ext  # [arcsec]
@@ -138,12 +137,14 @@ which are meant to feed into the `extent` parameter of `matplotlib.pyplot.imshow
 
 +++
 
-## The {class}`mpol.gridding.Gridder` object
+## Making images with {class}`~mpol.gridding.DirtyImager`
+
+Those familiar with radio astronomy will be familiar with the idea of "gridding" loose visibilities to a Cartesian $u,v$ grid. MPoL has two classes that "grid" visibilities: {class}`mpol.gridding.DirtyImager` and {class}`mpol.gridding.DataAverager`. Their internals may be similar, but they serve different purposes. First, let's look at how we can use the {class}`mpol.gridding.DirtyImager` to make diagnostic images using the inverse Fast Fourier Transform, frequently called the "dirty image" by radio astronomers.
 
-The purpose of the gridder is to take in loose visibility data (as from an ALMA observation) and average it to cells defined by the {class}`~mpol.coordinates.GridCoords` object. We can instantiate a {class}`~mpol.gridding.Gridder` object by
+We can instantiate a {class}`~mpol.gridding.DirtyImager` object by
 
 ```{code-cell}
-gridder = gridding.Gridder(
+imager = gridding.DirtyImager(
     coords=coords,
     uu=uu,
     vv=vv,
@@ -153,10 +154,10 @@ gridder = gridding.Gridder(
 )
 ```
 
-Instantiating the {class}`~mpol.gridding.Gridder` object attaches the {class}`~mpol.coordinates.GridCoords`  object and the loose visibilities. There is also a convenience method to create the {class}`~mpol.coordinates.GridCoords` and {class}`~mpol.gridding.Gridder` object in one shot by
+Instantiating the {class}`~mpol.gridding.DirtyImager` object attaches the {class}`~mpol.coordinates.GridCoords`  object and the loose visibilities. There is also a convenience method to create the {class}`~mpol.coordinates.GridCoords` and {class}`~mpol.gridding.DirtyImager` object in one shot by
 
 ```{code-cell}
-gridder = gridding.Gridder.from_image_properties(
+imager = gridding.DirtyImager.from_image_properties(
     cell_size=0.005,  # [arcsec]
     npix=800,
     uu=uu,
@@ -169,18 +170,14 @@ gridder = gridding.Gridder.from_image_properties(
 
 if you don't want to specify your {class}`~mpol.coordinates.GridCoords` object separately.
 
-+++
-
-## Making a diagnostic "dirty image" 
-
 As we saw, the raw visibility dataset is a set of complex-valued Fourier samples. Our objective is to make images of the sky-brightness distribution and do astrophysics. We'll cover how to do this with MPoL and RML techniques in later tutorials, but it is possible to get a rough idea of the sky brightness by calculating the inverse Fourier transform of the visibility values.
 
-To do this, you can call the {meth}`mpol.gridding.Gridder.get_dirty_image` method on your {class}`~mpol.gridding.Gridder` object. This routine will average, or 'grid', the loose visibilities to the Fourier grid defined by {class}`~mpol.coordinates.GridCoords` and then calculate the diagnostic dirty image and dirty beam cubes that correspond to the Fourier transform of the gridded visibilities.
+To do this, you can call the {meth}`mpol.gridding.DirtyImager.get_dirty_image` method on your {class}`~mpol.gridding.DirtyImager` object. This routine will average, or 'grid', the loose visibilities to the Fourier grid defined by {class}`~mpol.coordinates.GridCoords` and then calculate the diagnostic dirty image and dirty beam cubes that correspond to the Fourier transform of the gridded visibilities.
 
 There are several different schemes by which to do the averaging, each of which will deliver different image plane resolutions (defined by the size of the PSF or dirty beam) and thermal noise properties. MPoL implements 'uniform', 'natural', and 'briggs' robust weighting. For more information on the difference between these schemes, see the [CASA documentation](https://casa.nrao.edu/casadocs-devel/stable/imaging/synthesis-imaging/data-weighting) or Chapter 3 of Daniel Briggs' [Ph.D. thesis](http://www.aoc.nrao.edu/dissertations/dbriggs/).
 
 ```{code-cell}
-img, beam = gridder.get_dirty_image(weighting="briggs", robust=0.0)
+img, beam = imager.get_dirty_image(weighting="briggs", robust=0.0)
 ```
 
 Note that these are three dimensional image cubes with the same `nchan` as the input visibility data.
@@ -198,7 +195,7 @@ N.B. that the intensity units of the dirty image are technically undefined. The
 
 ```{code-cell}
 chan = 4
-kw = {"origin": "lower", "interpolation": "none", "extent": gridder.coords.img_ext}
+kw = {"origin": "lower", "interpolation": "none", "extent": imager.coords.img_ext}
 fig, ax = plt.subplots(ncols=2, figsize=(6.0, 4))
 ax[0].imshow(beam[chan], **kw)
 ax[0].set_title("beam")
@@ -212,21 +209,40 @@ fig.subplots_adjust(left=0.14, right=0.90, wspace=0.35, bottom=0.15, top=0.9)
 
 If you were working with this measurement set in CASA, it's a good idea to compare the dirty image produced here to the dirty image from CASA (i.e., produced by `tclean` with zero CLEAN iterations). You should confirm that these two dirty images look very similar (i.e., nearly but most likely not quite to numerical precision) before moving on to regularized maximum imaging. If your image appears upside down or mirrored, check whether you converted your visibility data from the CASA baseline convention to the regular TMS baseline convention by complex-conjugating your visibilities.
 
-+++
+
+## Averaging and exporting data with {class}`~mpol.gridding.DataAverager`
+
+As we saw at the beginning of this tutorial, an ALMA dataset may easily contain 1/3 million or more individual visibility measurements, which can present a computational burden for some imaging routines. Just like many noisy data points can be "binned" into a set of fewer higher signal to noise points (for example, as with a lightcurve of a transiting exoplanet), so too can visibility data points be averaged down. 
+
+To do this, you can instantiate a {class}`~mpol.gridding.DataAverager` object and then call the {meth}`mpol.gridding.DataAverager.to_pytorch_dataset` method. This routine will average, or 'grid', the loose visibilities to the Fourier grid defined by {class}`~mpol.coordinates.GridCoords` and then export the dataset as a {class}`mpol.datasets.GriddedDataset` object.
+
+```{code-cell}
+averager = gridding.DataAverager(
+    coords=coords,
+    uu=uu,
+    vv=vv,
+    weight=weight,
+    data_re=data_re,
+    data_im=data_im,
+)
+
+dset = averager.to_pytorch_dataset()
+```
+
 
 ## Checking data weights
-When working with real data, it is possible that the statistical uncertainties---conveyed by the weights---were [not correctly calibrated by certain CASA versions](https://mpol-dev.github.io/visread/tutorials/rescale_AS209_weights.html). For dirty and CLEAN imaging purposes, it's OK if the weights are not correctly scaled so long as their *relative* scalings are correct (to each other). For forward-modeling and RML imaging, it's important that the weights are correctly scaled in an absolute sense. To alert the user to the possibility that their weights may be incorrectly calibrated, the dirty imaging routines will raise a ``RuntimeWarning`` if the weights are incorrectly scaled. Even though the weights are incorrect, the dirty image may still be valid---hence why these routines issue a warning instead of an error.
+When working with real data, it is possible that the statistical uncertainties---conveyed by the weights---were [not correctly calibrated by certain CASA versions](https://mpol-dev.github.io/visread/tutorials/rescale_AS209_weights.html). For dirty and CLEAN imaging purposes, it's OK if the weights are not correctly scaled so long as their *relative* scalings are correct (to each other). For forward-modeling and RML imaging, it's important that the weights are correctly scaled in an absolute sense. To alert the user to the possibility that their weights may be incorrectly calibrated, the routines internal to {class}`~mpol.gridding.DirtyImager` will raise a ``RuntimeWarning`` if the weights are incorrectly scaled. Even though the weights are incorrect, the user image may still want the dirty image---hence why these routines issue a warning instead of an error.
 
 ```
-  img, beam = gridder.get_dirty_image(
+  img, beam = imager.get_dirty_image(
         weighting="uniform", check_visibility_scatter=True, max_scatter=1.2
   )
 ```
 
-However, if the user goes to export the gridded visibilities as a PyTorch dataset for RML imaging, incorrectly scaled weights will raise a ``RuntimeError``. RML images and forward modeling inferences will be compromised if the weights are not statistically valid.
+However, if the user goes to export the gridded visibilities as a PyTorch dataset for RML imaging using {class}`~mpol.gridding.DataAverager`, incorrectly scaled weights will raise a ``RuntimeError``. RML images and forward modeling inferences will be compromised if the weights are not statistically valid.
 
 The sensitivity of the export routines can be adjusted by changing the ``max_scatter`` keyword. Scatter checking can be disabled by setting ``check_visibility_scatter=False``, but is not recommended unless you are trying to debug things.
 
 ```
-  dset = gridder.to_pytorch_dataset(check_visibility_scatter=True, max_scatter=1.2)
+  dset = averager.to_pytorch_dataset(check_visibility_scatter=True, max_scatter=1.2)
 ```