diff --git a/autolens/__init__.py b/autolens/__init__.py
index 004171990..255e468eb 100644
--- a/autolens/__init__.py
+++ b/autolens/__init__.py
@@ -113,6 +113,11 @@
 from . import util
 
 from autoconf import conf
+from autoconf.fitsable import ndarray_via_hdu_from
+from autoconf.fitsable import ndarray_via_fits_from
+from autoconf.fitsable import header_obj_from
+from autoconf.fitsable import output_to_fits
+from autoconf.fitsable import hdu_list_for_output_from
 
 conf.instance.register(__file__)
 
diff --git a/autolens/aggregator/__init__.py b/autolens/aggregator/__init__.py
index b34e80ea5..1f156b3c0 100644
--- a/autolens/aggregator/__init__.py
+++ b/autolens/aggregator/__init__.py
@@ -27,4 +27,4 @@
 from autolens.aggregator.subplot import SubplotFit as subplot_fit
 from autolens.aggregator.subplot import SubplotFitLog10 as subplot_fit_log10
 from autolens.aggregator.subplot import FITSTracer as fits_tracer
-from autolens.aggregator.subplot import FITSFit as fits_fits
+from autolens.aggregator.subplot import FITSFit as fits_fit
diff --git a/autolens/aggregator/fit_imaging.py b/autolens/aggregator/fit_imaging.py
index 97d98f105..cdf595a47 100644
--- a/autolens/aggregator/fit_imaging.py
+++ b/autolens/aggregator/fit_imaging.py
@@ -18,12 +18,13 @@ def _fit_imaging_from(
     settings_inversion: aa.SettingsInversion = None,
 ) -> List[FitImaging]:
     """
-    Returns a list of `FitImaging` object from a `PyAutoFit` sqlite database `Fit` object.
+    Returns a list of `FitImaging` object from a `PyAutoFit` loaded directory `Fit` or sqlite database `Fit` object.
 
-    The results of a model-fit can be stored in a sqlite database, including the following attributes of the fit:
+    The results of a model-fit can be loaded from hard-disk or stored in a sqlite database, including the following
+    attributes of the fit:
 
     - The imaging data, noise-map, PSF and settings as .fits files (e.g. `dataset/data.fits`).
-    - The mask used to mask the `Imaging` data structure in the fit (`dataset/mask.fits`).
+    - The mask used to mask the `Imaging` data structure in the fit (`dataset.fits[hdu=0]`).
     - The settings of inversions used by the fit (`dataset/settings_inversion.json`).
 
     Each individual attribute can be loaded from the database via the `fit.value()` method.
@@ -41,7 +42,8 @@ def _fit_imaging_from(
     Parameters
     ----------
     fit
-        A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry in a sqlite database.
+        A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry which has been loaded from
+        an output directory or from an sqlite database..
     instance
         A manual instance that overwrites the max log likelihood instance in fit (e.g. for drawing the instance
         randomly from the PDF).
@@ -87,41 +89,42 @@ def __init__(
         settings_inversion: Optional[aa.SettingsInversion] = None,
     ):
         """
-        Interfaces with an `PyAutoFit` aggregator object to create instances of `FitImaging` objects from the results
-        of a model-fit.
-
-        The results of a model-fit can be stored in a sqlite database, including the following attributes of the fit:
-
-        - The imaging data, noise-map, PSF and settings as .fits files (e.g. `dataset/data.fits`).
-        - The mask used to mask the `Imaging` data structure in the fit (`dataset/mask.fits`).
-        - The settings of inversions used by the fit (`dataset/settings_inversion.json`).
-
-        The `aggregator` contains the path to each of these files, and they can be loaded individually. This class
-        can load them all at once and create an `FitImaging` object via the `_fit_imaging_from` method.
-
-        This class's methods returns generators which create the instances of the `FitImaging` objects. This ensures
-        that large sets of results can be efficiently loaded from the hard-disk and do not require storing all
-        `FitImaging` instances in the memory at once.
-
-        For example, if the `aggregator` contains 3 model-fits, this class can be used to create a generator which
-        creates instances of the corresponding 3 `FitImaging` objects.
-
-        If multiple `FitImaging` objects were fitted simultaneously via analysis summing, the `fit.child_values()` method
-        is instead used to load lists of the data, noise-map, PSF and mask and combine them into a list of
-        `FitImaging` objects.
-
-        This can be done manually, but this object provides a more concise API.
-
-        Parameters
-        ----------
-        aggregator
-            A `PyAutoFit` aggregator object which can load the results of model-fits.
-        settings_inversion
-            Optionally overwrite the `SettingsInversion` of the `Inversion` object that is created from the fit.
-        use_preloaded_grid
-            Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This
-            grid may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is
-            used as the fit.
+            Interfaces with an `PyAutoFit` aggregator object to create instances of `FitImaging` objects from the results
+            of a model-fit.
+
+            The results of a model-fit can be loaded from hard-disk or stored in a sqlite database, including the following
+        attributes of the fit:
+
+            - The imaging data, noise-map, PSF and settings as .fits files (e.g. `dataset/data.fits`).
+            - The mask used to mask the `Imaging` data structure in the fit (`dataset.fits[hdu=0]`).
+            - The settings of inversions used by the fit (`dataset/settings_inversion.json`).
+
+            The `aggregator` contains the path to each of these files, and they can be loaded individually. This class
+            can load them all at once and create an `FitImaging` object via the `_fit_imaging_from` method.
+
+            This class's methods returns generators which create the instances of the `FitImaging` objects. This ensures
+            that large sets of results can be efficiently loaded from the hard-disk and do not require storing all
+            `FitImaging` instances in the memory at once.
+
+            For example, if the `aggregator` contains 3 model-fits, this class can be used to create a generator which
+            creates instances of the corresponding 3 `FitImaging` objects.
+
+            If multiple `FitImaging` objects were fitted simultaneously via analysis summing, the `fit.child_values()` method
+            is instead used to load lists of the data, noise-map, PSF and mask and combine them into a list of
+            `FitImaging` objects.
+
+            This can be done manually, but this object provides a more concise API.
+
+            Parameters
+            ----------
+            aggregator
+                A `PyAutoFit` aggregator object which can load the results of model-fits.
+            settings_inversion
+                Optionally overwrite the `SettingsInversion` of the `Inversion` object that is created from the fit.
+            use_preloaded_grid
+                Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This
+                grid may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is
+                used as the fit.
         """
         super().__init__(aggregator=aggregator)
 
@@ -138,7 +141,8 @@ def object_via_gen_from(
         Parameters
         ----------
         fit
-            A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry in a sqlite database.
+            A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry which has been loaded from
+        an output directory or from an sqlite database..
         instance
             A manual instance that overwrites the max log likelihood instance in fit (e.g. for drawing the instance
             randomly from the PDF).
diff --git a/autolens/aggregator/fit_interferometer.py b/autolens/aggregator/fit_interferometer.py
index f0baf2729..bb0910759 100644
--- a/autolens/aggregator/fit_interferometer.py
+++ b/autolens/aggregator/fit_interferometer.py
@@ -15,13 +15,13 @@
 def _fit_interferometer_from(
     fit: af.Fit,
     instance: Optional[af.ModelInstance] = None,
-    real_space_mask: Optional[aa.Mask2D] = None,
     settings_inversion: aa.SettingsInversion = None,
 ) -> List[FitInterferometer]:
     """
-    Returns a list of `FitInterferometer` objects from a `PyAutoFit` sqlite database `Fit` object.
+    Returns a list of `FitInterferometer` objects from a `PyAutoFit` loaded directory `Fit` or sqlite database `Fit` object.
 
-    The results of a model-fit can be stored in a sqlite database, including the following attributes of the fit:
+    The results of a model-fit can be loaded from hard-disk or stored in a sqlite database, including the following
+    attributes of the fit:
 
     - The interferometer data, noise-map, uv-wavelengths and settings as .fits files (e.g. `dataset/data.fits`).
     - The real space mask defining the grid of the interferometer for the FFT (`dataset/real_space_mask.fits`).
@@ -43,7 +43,8 @@ def _fit_interferometer_from(
     Parameters
     ----------
     fit
-        A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry in a sqlite database.
+        A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry which has been loaded from
+        an output directory or from an sqlite database..
     instance
         A manual instance that overwrites the max log likelihood instance in fit (e.g. for drawing the instance
         randomly from the PDF).
@@ -52,7 +53,6 @@ def _fit_interferometer_from(
     """
     dataset_list = _interferometer_from(
         fit=fit,
-        real_space_mask=real_space_mask,
     )
     tracer_list = _tracer_from(fit=fit, instance=instance)
     dataset_model_list = _dataset_model_from(fit=fit, instance=instance)
@@ -87,45 +87,44 @@ def __init__(
         self,
         aggregator: af.Aggregator,
         settings_inversion: Optional[aa.SettingsInversion] = None,
-        real_space_mask: Optional[aa.Mask2D] = None,
     ):
         """
-        Interfaces with an `PyAutoFit` aggregator object to create instances of `FitInterferometer` objects from the
-        results of a model-fit.
+            Interfaces with an `PyAutoFit` aggregator object to create instances of `FitInterferometer` objects from the
+            results of a model-fit.
 
-        The results of a model-fit can be stored in a sqlite database, including the following attributes of the fit:
+            The results of a model-fit can be loaded from hard-disk or stored in a sqlite database, including the following
+        attributes of the fit:
 
-        - The interferometer data, noise-map, uv-wavelengths and settings as .fits files (e.g. `dataset/data.fits`).
-        - The real space mask defining the grid of the interferometer for the FFT (`dataset/real_space_mask.fits`).
-        - The settings of inversions used by the fit (`dataset/settings_inversion.json`).
+            - The interferometer data, noise-map, uv-wavelengths and settings as .fits files (e.g. `dataset/data.fits`).
+            - The real space mask defining the grid of the interferometer for the FFT (`dataset/real_space_mask.fits`).
+            - The settings of inversions used by the fit (`dataset/settings_inversion.json`).
 
-        The `aggregator` contains the path to each of these files, and they can be loaded individually. This class
-        can load them all at once and create an `FitInterferometer` object via the `_fit_interferometer_from` method.
+            The `aggregator` contains the path to each of these files, and they can be loaded individually. This class
+            can load them all at once and create an `FitInterferometer` object via the `_fit_interferometer_from` method.
 
-        This class's methods returns generators which create the instances of the `FitInterferometer` objects. This ensures
-        that large sets of results can be efficiently loaded from the hard-disk and do not require storing all
-        `FitInterferometer` instances in the memory at once.
+            This class's methods returns generators which create the instances of the `FitInterferometer` objects. This ensures
+            that large sets of results can be efficiently loaded from the hard-disk and do not require storing all
+            `FitInterferometer` instances in the memory at once.
 
-        For example, if the `aggregator` contains 3 model-fits, this class can be used to create a generator which
-        creates instances of the corresponding 3 `FitInterferometer` objects.
+            For example, if the `aggregator` contains 3 model-fits, this class can be used to create a generator which
+            creates instances of the corresponding 3 `FitInterferometer` objects.
 
-        This can be done manually, but this object provides a more concise API.
+            This can be done manually, but this object provides a more concise API.
 
-        Parameters
-        ----------
-        aggregator
-            A `PyAutoFit` aggregator object which can load the results of model-fits.
-        settings_inversion
-            Optionally overwrite the `SettingsInversion` of the `Inversion` object that is created from the fit.
-        use_preloaded_grid
-            Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This
-            grid may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is
-            used as the fit.
+            Parameters
+            ----------
+            aggregator
+                A `PyAutoFit` aggregator object which can load the results of model-fits.
+            settings_inversion
+                Optionally overwrite the `SettingsInversion` of the `Inversion` object that is created from the fit.
+            use_preloaded_grid
+                Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This
+                grid may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is
+                used as the fit.
         """
         super().__init__(aggregator=aggregator)
 
         self.settings_inversion = settings_inversion
-        self.real_space_mask = real_space_mask
 
     def object_via_gen_from(
         self, fit, instance: Optional[af.ModelInstance] = None
@@ -138,7 +137,8 @@ def object_via_gen_from(
             Parameters
             ----------
             fit
-                A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry in a sqlite database.
+                A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry which has been loaded from
+        an output directory or from an sqlite database..
         instance
             A manual instance that overwrites the max log likelihood instance in fit (e.g. for drawing the instance
             randomly from the PDF).
diff --git a/autolens/aggregator/tracer.py b/autolens/aggregator/tracer.py
index a0d1efb3a..e6febfe1f 100644
--- a/autolens/aggregator/tracer.py
+++ b/autolens/aggregator/tracer.py
@@ -12,9 +12,10 @@ def _tracer_from(
     fit: af.Fit, instance: Optional[af.ModelInstance] = None
 ) -> List[Tracer]:
     """
-    Returns a list of `Tracer` objects from a `PyAutoFit` sqlite database `Fit` object.
+    Returns a list of `Tracer` objects from a `PyAutoFit` loaded directory `Fit` or sqlite database `Fit` object.
 
-    The results of a model-fit can be stored in a sqlite database, including the following attributes of the fit:
+    The results of a model-fit can be loaded from hard-disk or stored in a sqlite database, including the following
+    attributes of the fit:
 
     - The model and its best fit parameters (e.g. `model.json`).
     - The adapt images associated with adaptive galaxy features (`adapt` folder).
@@ -31,7 +32,8 @@ def _tracer_from(
     Parameters
     ----------
     fit
-        A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry in a sqlite database.
+        A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry which has been loaded from
+        an output directory or from an sqlite database..
     instance
         A manual instance that overwrites the max log likelihood instance in fit (e.g. for drawing the instance
         randomly from the PDF).
@@ -79,7 +81,8 @@ class TracerAgg(af.AggBase):
     Interfaces with an `PyAutoFit` aggregator object to create instances of `Tracer` objects from the results
     of a model-fit.
 
-    The results of a model-fit can be stored in a sqlite database, including the following attributes of the fit:
+    The results of a model-fit can be loaded from hard-disk or stored in a sqlite database, including the following
+    attributes of the fit:
 
     - The model and its best fit parameters (e.g. `model.json`).
     - The adapt images associated with adaptive galaxy features (`adapt` folder).
@@ -117,7 +120,8 @@ def object_via_gen_from(
         Parameters
         ----------
         fit
-            A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry in a sqlite database.
+            A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry which has been loaded from
+        an output directory or from an sqlite database..
         galaxies
             A list of galaxies corresponding to a sample of a non-linear search and model-fit.
         """
diff --git a/autolens/analysis/plotter_interface.py b/autolens/analysis/plotter_interface.py
index e5e1b0cfb..9210d5dd6 100644
--- a/autolens/analysis/plotter_interface.py
+++ b/autolens/analysis/plotter_interface.py
@@ -1,4 +1,4 @@
-from os import path
+from autoconf.fitsable import hdu_list_for_output_from
 
 import autoarray as aa
 import autogalaxy.plot as aplt
@@ -65,25 +65,26 @@ def should_plot(name):
             tracer_plotter.subplot_galaxies_images()
 
         if should_plot("fits_tracer"):
-            number_plots = 4
-
-            multi_plotter = aplt.MultiFigurePlotter(
-                plotter_list=[tracer_plotter] * number_plots
-            )
-
-            multi_plotter.output_to_fits(
-                func_name_list=["figures_2d"] * number_plots,
-                figure_name_list=[
+            hdu_list = hdu_list_for_output_from(
+                values_list=[
+                    grid.mask.astype("float"),
+                    tracer.convergence_2d_from(grid=grid).native,
+                    tracer.potential_2d_from(grid=grid).native,
+                    tracer.deflections_yx_2d_from(grid=grid).native[:, :, 0],
+                    tracer.deflections_yx_2d_from(grid=grid).native[:, :, 1],
+                ],
+                ext_name_list=[
+                    "mask",
                     "convergence",
                     "potential",
                     "deflections_y",
                     "deflections_x",
                 ],
-                tag_list=["convergence", "potential", "deflections_y", "deflections_x"],
-                filename="tracer",
-                remove_fits_first=True,
+                header_dict=grid.mask.header_dict,
             )
 
+            hdu_list.writeto(self.image_path / "tracer.fits", overwrite=True)
+
     def image_with_positions(self, image: aa.Array2D, positions: aa.Grid2DIrregular):
         """
         Visualizes the positions of a model-fit, where these positions are used to resample lens models where
diff --git a/autolens/imaging/model/analysis.py b/autolens/imaging/model/analysis.py
index 64288c2c5..24a2ff8a1 100644
--- a/autolens/imaging/model/analysis.py
+++ b/autolens/imaging/model/analysis.py
@@ -169,22 +169,23 @@ def save_attributes(self, paths: af.DirectoryPaths):
         Before the non-linear search begins, this routine saves attributes of the `Analysis` object to the `files`
         folder such that they can be loaded after the analysis using PyAutoFit's database and aggregator tools.
 
-        For this analysis, it uses the `AnalysisDataset` object's method to output the following:
 
-        - The dataset's data.
-        - The dataset's noise-map.
-        - The settings associated with the dataset.
-        - The settings associated with the inversion.
-        - The settings associated with the pixelization.
-        - The Cosmology.
-        - The adapt image's model image and galaxy images, if used.
+         For this analysis, it uses the `AnalysisDataset` object's method to output the following:
 
-        This function also outputs attributes specific to an imaging dataset:
+         - The settings associated with the inversion.
+         - The settings associated with the pixelization.
+         - The Cosmology.
+         - The adapt image's model image and galaxy images, as `adapt_images.fits`, if used.
+
+        This function also outputs attributes specific to lens modeling:
 
-        - Its PSF.
-        - Its mask.
         - The positions of the brightest pixels in the lensed source which are used to discard mass models.
 
+        The following .fits files are also output via the plotter interface:
+
+        - The mask applied to the dataset, in the `PrimaryHDU` of `dataset.fits`.
+        - The imaging dataset as `dataset.fits` (data / noise-map / psf / over sampler / etc.).
+
         It is common for these attributes to be loaded by many of the template aggregator functions given in the
         `aggregator` modules. For example, when using the database tools to perform a fit, the default behaviour is for
         the dataset, settings and other attributes necessary to perform the fit to be loaded via the pickle files
diff --git a/autolens/imaging/model/plotter_interface.py b/autolens/imaging/model/plotter_interface.py
index fbf023b3b..b307273e5 100644
--- a/autolens/imaging/model/plotter_interface.py
+++ b/autolens/imaging/model/plotter_interface.py
@@ -77,7 +77,7 @@ def should_plot(name):
             except IndexError:
                 pass
 
-        fits_to_fits(should_plot=should_plot, fit=fit, mat_plot_2d=mat_plot_2d, fit_plotter_cls=FitImagingPlotter)
+        fits_to_fits(should_plot=should_plot, image_path=self.image_path, fit=fit)
 
     def fit_imaging_combined(self, fit_list: List[FitImaging]):
         """
diff --git a/autolens/imaging/simulator.py b/autolens/imaging/simulator.py
index 7d68d2007..2446b95f8 100644
--- a/autolens/imaging/simulator.py
+++ b/autolens/imaging/simulator.py
@@ -45,7 +45,11 @@ def via_tracer_from(self, tracer : Tracer, grid : aa.type.Grid2DLike) -> aa.Imag
             grid=grid, psf_shape_2d=self.psf.shape_native
         )
 
-        dataset = self.via_image_from(image=image)
+        over_sample_size = grid.over_sample_size.resized_from(
+            new_shape=image.shape_native, mask_pad_value=1
+        )
+
+        dataset = self.via_image_from(image=image, over_sample_size=over_sample_size)
 
         return dataset.trimmed_after_convolution_from(
             kernel_shape=self.psf.shape_native
diff --git a/autolens/interferometer/model/analysis.py b/autolens/interferometer/model/analysis.py
index 141a941db..3dec88777 100644
--- a/autolens/interferometer/model/analysis.py
+++ b/autolens/interferometer/model/analysis.py
@@ -223,25 +223,25 @@ def fit_from(
 
     def save_attributes(self, paths: af.DirectoryPaths):
         """
-         Before the non-linear search begins, this routine saves attributes of the `Analysis` object to the `files`
-         folder such that they can be loaded after the analysis using PyAutoFit's database and aggregator tools.
+         Before the model-fit begins, this routine saves attributes of the `Analysis` object to the `files` folder
+         such that they can be loaded after the analysis using PyAutoFit's database and aggregator tools.
 
          For this analysis, it uses the `AnalysisDataset` object's method to output the following:
 
-         - The dataset's data.
-         - The dataset's noise-map.
-         - The settings associated with the dataset.
          - The settings associated with the inversion.
          - The settings associated with the pixelization.
          - The Cosmology.
-         - The adapt image's model image and galaxy images, if used.
+         - The adapt image's model image and galaxy images, as `adapt_images.fits`, if used.
 
-         This function also outputs attributes specific to an imaging dataset:
+         This function also outputs attributes specific to lens modeling:
 
-        - Its uv-wavelengths
-        - Its real space mask.
         - The positions of the brightest pixels in the lensed source which are used to discard mass models.
 
+        The following .fits files are also output via the plotter interface:
+
+        - The real space mask applied to the dataset, in the `PrimaryHDU` of `dataset.fits`.
+        - The interferometer dataset as `dataset.fits` (data / noise-map / uv_wavelengths).
+
          It is common for these attributes to be loaded by many of the template aggregator functions given in the
          `aggregator` modules. For example, when using the database tools to perform a fit, the default behaviour is for
          the dataset, settings and other attributes necessary to perform the fit to be loaded via the pickle files
diff --git a/autolens/interferometer/model/plotter_interface.py b/autolens/interferometer/model/plotter_interface.py
index ba5c03c06..827544d41 100644
--- a/autolens/interferometer/model/plotter_interface.py
+++ b/autolens/interferometer/model/plotter_interface.py
@@ -79,7 +79,6 @@ def should_plot(name):
 
         fits_to_fits(
             should_plot=should_plot,
+            image_path=self.image_path,
             fit=fit,
-            mat_plot_2d=mat_plot_2d,
-            fit_plotter_cls=FitInterferometerPlotter,
         )
diff --git a/test_autolens/aggregator/conftest.py b/test_autolens/aggregator/conftest.py
index 5fe1a6755..8e9419fb3 100644
--- a/test_autolens/aggregator/conftest.py
+++ b/test_autolens/aggregator/conftest.py
@@ -44,8 +44,11 @@ def aggregator_from(database_file, analysis, model, samples):
         samples=samples, result=al.m.MockResult(model=model, samples=samples)
     )
     search.paths = af.DirectoryPaths(path_prefix=database_file)
+
     search.fit(model=model, analysis=analysis)
 
+    analysis.visualize_before_fit(paths=search.paths, model=model)
+
     database_file = path.join(conf.instance.output_path, f"{database_file}.sqlite")
 
     agg = af.Aggregator.from_database(filename=database_file)
diff --git a/test_autolens/analysis/test_plotter_interface.py b/test_autolens/analysis/test_plotter_interface.py
index 27506324a..9f6142615 100644
--- a/test_autolens/analysis/test_plotter_interface.py
+++ b/test_autolens/analysis/test_plotter_interface.py
@@ -29,11 +29,11 @@ def test__tracer(
 
     assert path.join(plot_path, "subplot_galaxies_images.png") in plot_patch.paths
 
-    image = al.util.array_2d.numpy_array_2d_via_fits_from(
+    image = al.ndarray_via_fits_from(
         file_path=path.join(plot_path, "tracer.fits"), hdu=0
     )
 
-    assert image.shape == (5, 5)
+    assert image.shape == (7, 7)
 
 
 def test__image_with_positions(
diff --git a/test_autolens/imaging/model/test_plotter_interface_imaging.py b/test_autolens/imaging/model/test_plotter_interface_imaging.py
index c3cb38ea4..6ee0d0cc7 100644
--- a/test_autolens/imaging/model/test_plotter_interface_imaging.py
+++ b/test_autolens/imaging/model/test_plotter_interface_imaging.py
@@ -30,17 +30,17 @@ def test__fit_imaging(
     assert path.join(plot_path, "subplot_fit.png") in plot_patch.paths
     assert path.join(plot_path, "subplot_fit_log10.png") in plot_patch.paths
 
-    image = al.util.array_2d.numpy_array_2d_via_fits_from(
+    image = al.ndarray_via_fits_from(
         file_path=path.join(plot_path, "fit.fits"), hdu=0
     )
 
-    assert image.shape == (5, 5)
+    assert image.shape == (7, 7)
 
-    image = al.util.array_2d.numpy_array_2d_via_fits_from(
+    image = al.ndarray_via_fits_from(
         file_path=path.join(plot_path, "model_galaxy_images.fits"), hdu=0
     )
 
-    assert image.shape == (5, 5)
+    assert image.shape == (7, 7)
 
 def test__fit_imaging_combined(
     fit_imaging_x2_plane_inversion_7x7, plot_path, plot_patch
diff --git a/test_autolens/interferometer/model/test_plotter_interface_interferometer.py b/test_autolens/interferometer/model/test_plotter_interface_interferometer.py
index 8a883c4fd..285535de0 100644
--- a/test_autolens/interferometer/model/test_plotter_interface_interferometer.py
+++ b/test_autolens/interferometer/model/test_plotter_interface_interferometer.py
@@ -32,19 +32,13 @@ def test__fit_interferometer(
     assert path.join(plot_path, "subplot_fit_real_space.png") in plot_patch.paths
     assert path.join(plot_path, "subplot_fit_dirty_images.png") in plot_patch.paths
 
-    # visibilities = ag.util.array_2d.numpy_array_2d_via_fits_from(
-    #     file_path=path.join(plot_path, "fit.fits"), hdu=0
-    # )
-    #
-    # assert visibilities.shape == (5, 5)
-
-    image = al.util.array_2d.numpy_array_2d_via_fits_from(
+    image = al.ndarray_via_fits_from(
         file_path=path.join(plot_path, "model_galaxy_images.fits"), hdu=0
     )
 
-    assert image.shape == (5, 5)
+    assert image.shape == (7, 7)
 
-    image = al.util.array_2d.numpy_array_2d_via_fits_from(
+    image = al.ndarray_via_fits_from(
         file_path=path.join(plot_path, "dirty_images.fits"), hdu=0
     )