From f847e0d74de8c32c5849951ab6a2255da0fa4769 Mon Sep 17 00:00:00 2001 From: Ian Czekala Date: Wed, 27 Dec 2023 21:37:16 +0000 Subject: [PATCH] Tests pass after removing passthrough. --- pyproject.toml | 10 +++--- src/mpol/fourier.py | 2 +- src/mpol/images.py | 79 +++++++++-------------------------------- src/mpol/precomposed.py | 18 +++++----- test/conftest.py | 14 ++++++-- test/connectors_test.py | 4 +-- test/fourier_test.py | 9 +++-- test/images_test.py | 8 +++-- test/losses_test.py | 48 +++++++++---------------- test/train_test_test.py | 2 +- 10 files changed, 74 insertions(+), 120 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index af5286e4..365a3bb1 100644 --- a/pyproject.toml +++ b/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 \ No newline at end of file diff --git a/src/mpol/fourier.py b/src/mpol/fourier.py index 706d3c67..ca1dbafd 100644 --- a/src/mpol/fourier.py +++ b/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] diff --git a/src/mpol/images.py b/src/mpol/images.py index 18133aa6..628fd55f 100644 --- a/src/mpol/images.py +++ b/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 diff --git a/src/mpol/precomposed.py b/src/mpol/precomposed.py index 35c2f6f6..227b3e3d 100644 --- a/src/mpol/precomposed.py +++ b/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. diff --git a/test/conftest.py b/test/conftest.py index d63d3377..b52b09cc 100644 --- a/test/conftest.py +++ b/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) diff --git a/test/connectors_test.py b/test/connectors_test.py index 5dcb801d..e0303983 100644 --- a/test/connectors_test.py +++ b/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())) diff --git a/test/fourier_test.py b/test/fourier_test.py index 39cb1fbb..40e91969 100644 --- a/test/fourier_test.py +++ b/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)) diff --git a/test/images_test.py b/test/images_test.py index eab22fca..b365c763 100644 --- a/test/images_test.py +++ b/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 diff --git a/test/losses_test.py b/test/losses_test.py index 908eb9e5..ca88a5c7 100644 --- a/test/losses_test.py +++ b/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( diff --git a/test/train_test_test.py b/test/train_test_test.py index 9826d9fc..04e6b66a 100644 --- a/test/train_test_test.py +++ b/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