Tests pass after removing passthrough.

pyproject.toml

@@ -91,10 +91,12 @@ ignore_missing_imports = true
 
 [[tool.mypy.overrides]]
 module = [
-    "MPol.constants",
-    "MPoL.losses",
+    "MPoL.constants",
+    "MPoL.coordinates",
+    "MPoL.datasets",
     # "MPoL.fourier", # once we remove get_vis_residuals
+    "MPoL.images",
+    "MPoL.losses"
     # "MPoL.utils", # once we fix check_baselines
-    "MPoL.datasets"
-]
+    ]
 disallow_untyped_defs = true

src/mpol/fourier.py

@@ -852,7 +852,7 @@ def get_vis_residuals(model, u_true, v_true, V_true, return_Vmod=False, channel=
     """
     nufft = NuFFT(coords=model.coords, nchan=model.nchan)
 
-    vis_model = nufft(model.icube().to("cpu"), u_true, v_true)  # TODO: remove 'to' call
+    vis_model = nufft(model.icube.cube.to("cpu"), u_true, v_true)  # TODO: remove 'to' call
     # convert to numpy, select channel
     vis_model = vis_model.detach().numpy()[channel]
 

src/mpol/images.py

@@ -61,9 +61,8 @@ def __init__(
         # The ``base_cube`` is already packed to make the Fourier transformation easier
         if base_cube is None:
             self.base_cube = nn.Parameter(
-                torch.full(
+                torch.zeros(
                     (self.nchan, self.coords.npix, self.coords.npix),
-                    fill_value=0.05,
                     requires_grad=True,
                     dtype=torch.double,
                 )
@@ -211,80 +210,36 @@ class ImageCube(nn.Module):
         the image `cell_size` and `npix`.
     nchan : int
         the number of channels in the base cube. Default = 1.
-    passthrough : bool
-        if `True`, assume ImageCube is just a layer as opposed
-        to parameter base.
-    cube : :class:torch.Tensor of :class:torch.double, of shape ``(nchan, npix, npix)``
-        a prepacked image cube to initialize the model with in units of
-        [:math:`\mathrm{Jy}\,\mathrm{arcsec}^{-2}`]. If None, assumes starting
-        ``cube`` is ``torch.zeros``. See :ref:`cube-orientation-label` for more
-        information on the expectations of the orientation of the input image.
     """
 
     def __init__(
         self,
         coords: GridCoords,
         nchan: int = 1,
-        passthrough: bool = False,
-        cube: torch.Tensor | None = None,
     ) -> None:
         super().__init__()
 
         self.coords = coords
         self.nchan = nchan
+        self.register_buffer("cube", None)
 
-        self.passthrough = passthrough
-
-        if not self.passthrough:
-            if cube is None:
-                self.cube : torch.nn.Parameter = nn.Parameter(
-                    torch.full(
-                        (self.nchan, self.coords.npix, self.coords.npix),
-                        fill_value=0.0,
-                        requires_grad=True,
-                        dtype=torch.double,
-                    )
-                )
-
-            else:
-                # We expect the user to supply a pre-packed base cube
-                # so that it's ready to go for the FFT
-                # We could apply this transformation for the user, but I think it will
-                # lead to less confusion if we make this transformation explicit
-                # for the user during the setup phase.
-                self.cube = nn.Parameter(cube)
-        else:
-            # ImageCube is working as a passthrough layer, so cube should
-            # only be provided as an arg to the forward method, not as
-            # an initialization argument
-            self.cube = None
-
-    def forward(self, cube: torch.Tensor | None = None) -> torch.Tensor:
+    def forward(self, cube: torch.Tensor) -> torch.Tensor:
         r"""
-        If the ImageCube object was initialized with ``passthrough=True``, the ``cube``
-        argument is required. ``forward`` essentially just passes this on as an identity
-        operation.
-
-        If the ImageCube object was initialized with ``passthrough=False``, the ``cube``
-        argument is not permitted, and ``forward`` passes on the stored
-        ``nn.Parameter`` cube as an identity operation.
-
-        Args:
-            cube (3D torch tensor of shape ``(nchan, npix, npix)``): only permitted if
-            the ImageCube object was initialized with ``passthrough=True``.
-
-        Returns: (3D torch.double tensor of shape ``(nchan, npix, npix)``) as identity
-        operation
+        Pass the cube through as an identity operation, storing the value to the 
+        internal buffer. After the cube has been passed through, convenience 
+        instance attributes like `sky_cube` and `flux` will reflect the updated cube.
+
+        Parameters
+        ----------
+        cube : :class:`torch.Tensor` of type :class:`torch.double` 
+            3D torch tensor of shape ``(nchan, npix, npix)``) in 'packed' format
+
+        Returns
+        -------
+        :class:`torch.Tensor` of :class:`torch.double` type
+            tensor of shape ``(nchan, npix, npix)``), same as `cube`
         """
-
-        if cube is not None:
-            assert (
-                self.passthrough
-            ), "ImageCube.passthrough must be True if supplying cube."
-            self.cube = cube
-
-        if not self.passthrough:
-            assert cube is None, "Do not supply cube if ImageCube.passthrough == False."
+        self.cube = cube
 
         return self.cube
 

src/mpol/precomposed.py

@@ -13,10 +13,10 @@ class SimpleNet(torch.nn.Module):
     Args:
         cell_size (float): the width of a pixel [arcseconds]
         npix (int): the number of pixels per image side
-        coords (GridCoords): an object already instantiated from the GridCoords class. 
+        coords (GridCoords): an object already instantiated from the GridCoords class.
             If providing this, cannot provide ``cell_size`` or ``npix``.
         nchan (int): the number of channels in the base cube. Default = 1.
-        base_cube : a pre-packed base cube to initialize the model with. If 
+        base_cube : a pre-packed base cube to initialize the model with. If
             None, assumes ``torch.zeros``.
 
     After the object is initialized, instance variables can be accessed, for example
@@ -25,10 +25,10 @@ class SimpleNet(torch.nn.Module):
     :ivar icube: the :class:`~mpol.images.ImageCube` instance
     :ivar fcube: the :class:`~mpol.fourier.FourierCube` instance
 
-    For example, you'll likely want to access the ``self.icube.sky_model`` 
+    For example, you'll likely want to access the ``self.icube.sky_model``
     at some point.
 
-    The idea is that :class:`~mpol.precomposed.SimpleNet` can save you some keystrokes 
+    The idea is that :class:`~mpol.precomposed.SimpleNet` can save you some keystrokes
     composing models by connecting the most commonly used layers together.
 
     .. mermaid:: _static/mmd/src/SimpleNet.mmd
@@ -52,16 +52,14 @@ def __init__(
 
         self.conv_layer = images.HannConvCube(nchan=self.nchan)
 
-        self.icube = images.ImageCube(
-            coords=self.coords, nchan=self.nchan, passthrough=True
-        )
+        self.icube = images.ImageCube(coords=self.coords, nchan=self.nchan)
         self.fcube = fourier.FourierCube(coords=self.coords)
 
     def forward(self):
         r"""
-        Feed forward to calculate the model visibilities. In this step, a 
-        :class:`~mpol.images.BaseCube` is fed to a :class:`~mpol.images.HannConvCube` 
-        is fed to a :class:`~mpol.images.ImageCube` is fed to a 
+        Feed forward to calculate the model visibilities. In this step, a
+        :class:`~mpol.images.BaseCube` is fed to a :class:`~mpol.images.HannConvCube`
+        is fed to a :class:`~mpol.images.ImageCube` is fed to a
         :class:`~mpol.fourier.FourierCube` to produce the visibility cube.
 
         Returns: 1D complex torch tensor of model visibilities.

test/conftest.py

@@ -149,7 +149,10 @@ def mock_1d_image_model(mock_1d_archive):
     # 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)
+    bcube = images.BaseCube(coords=coords,nchan=1,base_cube=packed_tensor,pixel_mapping=lambda x: x)
+    cube_true = images.ImageCube(coords=coords, nchan=1)
+    # register cube to buffer inside cube_true.cube
+    cube_true(bcube())
 
     return rtrue, itrue, cube_true, xmax, ymax, geom
 
@@ -176,8 +179,13 @@ def mock_1d_vis_model(mock_1d_archive):
     # 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)
-
+    bcube = images.BaseCube(coords=coords,nchan=1, base_cube=packed_tensor, pixel_mapping=lambda x:x)
+    cube_true = images.ImageCube(coords=coords, nchan=1)
+
+    # register image 
+    cube_true(bcube())
+
+
     # create a FourierCube
     fcube_true = fourier.FourierCube(coords=coords)    
 

test/connectors_test.py

@@ -27,7 +27,7 @@ def test_index(coords, dataset):
     basecube = images.BaseCube(coords=coords, nchan=nchan, base_cube=base_cube)
 
     # try passing through ImageLayer
-    imagecube = images.ImageCube(coords=coords, nchan=nchan, passthrough=True)
+    imagecube = images.ImageCube(coords=coords, nchan=nchan)
 
     # produce dense model visibility cube
     modelVisibilityCube = flayer(imagecube(basecube()))
@@ -42,7 +42,7 @@ def test_connector_grad(coords, dataset):
     flayer = fourier.FourierCube(coords=coords)
     nchan = dataset.nchan
     basecube = images.BaseCube(coords=coords, nchan=nchan)
-    imagecube = images.ImageCube(coords=coords, nchan=nchan, passthrough=True)
+    imagecube = images.ImageCube(coords=coords, nchan=nchan)
 
     # produce model visibilities
     modelVisibilityCube = flayer(imagecube(basecube()))

test/fourier_test.py

@@ -142,7 +142,8 @@ def test_predict_vis_nufft(coords, mock_visibility_data_cont):
 
     nchan = 10
 
-    # instantiate an ImageCube layer filled with zeros
+    # instantiate an BaseCube layer filled with zeros
+    basecube = images.BaseCube(coords=coords, nchan=nchan, pixel_mapping=lambda x: x)
     imagecube = images.ImageCube(coords=coords, nchan=nchan)
 
     # we have a multi-channel cube, but only sent single-channel uu and vv
@@ -151,7 +152,7 @@ def test_predict_vis_nufft(coords, mock_visibility_data_cont):
     layer = fourier.NuFFT(coords=coords, nchan=nchan)
 
     # predict the values of the cube at the u,v locations
-    output = layer(imagecube(), uu, vv)
+    output = layer(imagecube(basecube()), uu, vv)
 
     # make sure we got back the number of visibilities we expected
     assert output.shape == (nchan, len(uu))
@@ -172,6 +173,8 @@ def test_predict_vis_nufft_cached(coords, mock_visibility_data_cont):
     nchan = 10
 
     # instantiate an ImageCube layer filled with zeros
+    # instantiate an BaseCube layer filled with zeros
+    basecube = images.BaseCube(coords=coords, nchan=nchan, pixel_mapping=lambda x: x)
     imagecube = images.ImageCube(coords=coords, nchan=nchan)
 
     # we have a multi-channel cube, but sent only single-channel uu and vv
@@ -180,7 +183,7 @@ def test_predict_vis_nufft_cached(coords, mock_visibility_data_cont):
     layer = fourier.NuFFTCached(coords=coords, nchan=nchan, uu=uu, vv=vv)
 
     # predict the values of the cube at the u,v locations
-    output = layer(imagecube())
+    output = layer(imagecube(basecube()))
 
     # make sure we got back the number of visibilities we expected
     assert output.shape == (nchan, len(uu))

test/images_test.py

@@ -22,7 +22,7 @@ def test_basecube_grad():
 def test_imagecube_grad(coords):
     bcube = images.BaseCube(coords=coords)
     # try passing through ImageLayer
-    imagecube = images.ImageCube(coords=coords, passthrough=True)
+    imagecube = images.ImageCube(coords=coords)
 
     # send things through this layer
     loss = torch.sum(imagecube(bcube()))
@@ -36,7 +36,7 @@ def test_imagecube_tofits(coords, tmp_path):
     bcube = images.BaseCube(coords=coords)
 
     # try passing through ImageLayer
-    imagecube = images.ImageCube(coords=coords, passthrough=True)
+    imagecube = images.ImageCube(coords=coords)
 
     # sending the basecube through the imagecube
     imagecube(bcube())
@@ -88,7 +88,7 @@ def test_basecube_imagecube(coords, tmp_path):
     fig.savefig(tmp_path / "basecube_mapped.png", dpi=300)
 
     # try passing through ImageLayer
-    imagecube = images.ImageCube(coords=coords, nchan=nchan, passthrough=True)
+    imagecube = images.ImageCube(coords=coords, nchan=nchan)
 
     # send things through this layer
     imagecube(basecube())
@@ -173,5 +173,7 @@ def test_multi_chan_conv(coords, tmp_path):
 
 def test_image_flux(coords):
     nchan = 20
+    bcube = images.BaseCube(coords=coords, nchan=nchan)
     im = images.ImageCube(coords=coords, nchan=nchan)
+    im(bcube())
     assert im.flux.size()[0] == nchan

test/losses_test.py

@@ -5,43 +5,29 @@
 from mpol import fourier, images, losses, utils
 
 
-# create a fixture that has an image and produces loose and gridded model visibilities
+# create a fixture that returns nchan and an image
 @pytest.fixture
-def image_cube(mock_visibility_data, coords):
-    # Gaussian parameters
-    kw = {
-        "a": 1,
-        "delta_x": 0.02,  # arcsec
-        "delta_y": -0.01,
-        "sigma_x": 0.02,
-        "sigma_y": 0.01,
-        "Omega": 20,  # degrees
-    }
-
+def nchan_cube(mock_visibility_data, coords):
     uu, vv, weight, data_re, data_im = mock_visibility_data
     nchan = len(uu)
 
-    # evaluate the Gaussian over the sky-plane, as np array
-    img_packed = utils.sky_gaussian_arcsec(
-        coords.packed_x_centers_2D, coords.packed_y_centers_2D, **kw
+    # create a mock base image
+    basecube = images.BaseCube(
+        coords=coords,
+        nchan=nchan,
     )
-
-    # broadcast to (nchan, npix, npix)
-    img_packed_cube = np.broadcast_to(
-        img_packed, (nchan, coords.npix, coords.npix)
-    ).copy()
-    # convert img_packed to pytorch tensor
-    img_packed_tensor = torch.from_numpy(img_packed_cube)
     # insert into ImageCube layer
+    imagecube = images.ImageCube(coords=coords, nchan=nchan)
+    packed_cube = imagecube(basecube())
 
-    return images.ImageCube(coords=coords, nchan=nchan, cube=img_packed_tensor)
+    return nchan, packed_cube
 
 
 @pytest.fixture
-def loose_visibilities(mock_visibility_data, image_cube):
+def loose_visibilities(mock_visibility_data, coords, nchan_cube):
     # use the NuFFT to produce model visibilities
 
-    nchan = image_cube.nchan
+    nchan, packed_cube = nchan_cube
 
     # use the coil broadcasting ability
     chan = 4
@@ -50,16 +36,16 @@ def loose_visibilities(mock_visibility_data, image_cube):
     uu_chan = uu[chan]
     vv_chan = vv[chan]
 
-    nufft = fourier.NuFFT(coords=image_cube.coords, nchan=nchan)
-    return nufft(image_cube(), uu_chan, vv_chan)
+    nufft = fourier.NuFFT(coords=coords, nchan=nchan)
+    return nufft(packed_cube, uu_chan, vv_chan)
 
 
 @pytest.fixture
-def gridded_visibilities(image_cube):
-
+def gridded_visibilities(coords, nchan_cube):
+    nchan, packed_cube = nchan_cube
     # use the FourierCube to produce model visibilities
-    flayer = fourier.FourierCube(coords=image_cube.coords)
-    return flayer(image_cube())
+    flayer = fourier.FourierCube(coords=coords)
+    return flayer(packed_cube)
 
 
 def test_chi_squared_evaluation(

test/train_test_test.py

@@ -175,7 +175,7 @@ def test_train_to_dirty_image(coords, dataset, imager):
     train_to_dirty_image(model, imager, niter=10)
 
 
-def test_tensorboard(coords, dataset_cont, tmp_path):
+def test_tensorboard(coords, dataset_cont):
     # not using TrainTest class, 
     # set everything up to run on a single channel and then
     # test the writer function