Repository: TCDSolar/xrayvision (language: python)

Issue:
Support phase and image centre being different
Take a visibility with a phase and map centre of 0,0 and create and image between -5 to +5 in x and y the same image should be created with a phase centre of -5, -5, map centre of 0,0 created between 0 +11 and 0, +11 (if the visibilities are adjusted correctly)

It's the product of `xu` and `yv` that matters for the image/array values the centre 'coordinate' is need for coordinate aware transforms, plots etc


Discussion / hints:

Fix patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index d16b96e..6feceb9 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -27,6 +27,11 @@ repos:
         exclude: ".*(.fits|.fts|.fit|.header|.txt|tca.*|.json)$|^CITATION.rst$"
       - id: mixed-line-ending
         exclude: ".*(.fits|.fts|.fit|.header|.txt|tca.*)$"
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: "v1.10.0"
+    hooks:
+      - id: mypy
+        additional_dependencies: [ "types-setuptools" ]
   - repo: https://github.com/codespell-project/codespell
     rev: v2.2.6
     hooks:
diff --git a/changelog/58.breaking.rst b/changelog/58.breaking.rst
new file mode 100644
index 0000000..4840028
--- /dev/null
+++ b/changelog/58.breaking.rst
@@ -0,0 +1,2 @@
+Optional parameters are now keyword only for the :mod:`xrayvision.transform`, :mod:`xrayvision.imaging` and :mod:`xrayvision.visibility` modules.
+Remove ``natural`` keyword in favour of ``scheme`` keyword which can be either 'natural' or 'uniform'.
diff --git a/changelog/58.bugfix.rst b/changelog/58.bugfix.rst
new file mode 100644
index 0000000..325b4ce
--- /dev/null
+++ b/changelog/58.bugfix.rst
@@ -0,0 +1,1 @@
+Fix a bug where the x, y dimensions were not being treated consistently in :mod:`xrayvision.transform`.
diff --git a/changelog/58.docs.rst b/changelog/58.docs.rst
new file mode 100644
index 0000000..3499690
--- /dev/null
+++ b/changelog/58.docs.rst
@@ -0,0 +1,1 @@
+Add per module reference pages, switch to documenting types using type annotations.
diff --git a/docs/conf.py b/docs/conf.py
index 6f20ad1..28d1c4e 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -46,6 +46,9 @@
 # Add any paths that contain templates here, relative to this directory.
 # templates_path = ['_templates']
 
+# Set automodapi to generate files inside the generated directory
+automodapi_toctreedirnm = "generated/api"
+
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
@@ -63,7 +66,7 @@
 default_role = "obj"
 
 # Disable having a separate return type row
-napoleon_use_rtype = False
+napoleon_use_rtype = True
 
 # Disable google style docstrings
 napoleon_google_docstring = False
@@ -71,6 +74,9 @@
 # until sphinx-gallery / sphinx is fixed https://github.com/sphinx-doc/sphinx/issues/12300
 suppress_warnings = ["config.cache"]
 
+autodoc_typehints = "description"
+autoclass_content = "init"
+
 # -- Options for intersphinx extension ---------------------------------------
 
 # Example configuration for intersphinx: refer to the Python standard library.
diff --git a/docs/reference/clean.rst b/docs/reference/clean.rst
new file mode 100644
index 0000000..a3911f2
--- /dev/null
+++ b/docs/reference/clean.rst
@@ -0,0 +1,8 @@
+.. _clean:
+
+Clean ('xrayvision.clean')
+**************************
+
+The ``clean`` submodule contains clean imaging methods.
+
+.. automodapi:: xrayvision.clean
diff --git a/docs/reference/imaging.rst b/docs/reference/imaging.rst
new file mode 100644
index 0000000..959cb52
--- /dev/null
+++ b/docs/reference/imaging.rst
@@ -0,0 +1,8 @@
+.. _imaging:
+
+Imaging ('xrayvision.imaging')
+******************************
+
+The ``imaging`` submodule contains functions to make map and images from visibilities.
+
+.. automodapi:: xrayvision.imaging
diff --git a/docs/reference/index.rst b/docs/reference/index.rst
index cd0ca4e..389d57b 100644
--- a/docs/reference/index.rst
+++ b/docs/reference/index.rst
@@ -1,23 +1,15 @@
-.. _reference:
-
 *********
 Reference
 *********
 
-.. automodapi:: xrayvision.visibility
-
-.. automodapi:: xrayvision.imaging
-
-.. automodapi:: xrayvision.transform
-
-.. automodapi:: xrayvision.clean
-
-.. automodapi:: xrayvision.mem
-
-.. automodapi:: xrayvision.utils
-
-
 .. toctree::
-   :maxdepth: 2
+   :maxdepth: 1
+
+   clean
+   imaging
+   mem
+   transform
+   utils
+   visibility
 
    ../whatsnew/index
diff --git a/docs/reference/mem.rst b/docs/reference/mem.rst
new file mode 100644
index 0000000..c6141ca
--- /dev/null
+++ b/docs/reference/mem.rst
@@ -0,0 +1,8 @@
+.. _mem:
+
+MEM ('xrayvision.mem')
+**********************
+
+The ``mem`` submodule contains the Maximum Entropy methods.
+
+.. automodapi:: xrayvision.mem
diff --git a/docs/reference/transform.rst b/docs/reference/transform.rst
new file mode 100644
index 0000000..fa0a3bd
--- /dev/null
+++ b/docs/reference/transform.rst
@@ -0,0 +1,8 @@
+.. _transform:
+
+Transform ('xrayvision.transform')
+**********************************
+
+The ``transform`` submodule forward and reverse transforms.
+
+.. automodapi:: xrayvision.transform
diff --git a/docs/reference/utils.rst b/docs/reference/utils.rst
new file mode 100644
index 0000000..08f8df0
--- /dev/null
+++ b/docs/reference/utils.rst
@@ -0,0 +1,8 @@
+.. _utils:
+
+Utils ('xrayvision.utils')
+**************************
+
+The ``utils`` submodule contains utility functions.
+
+.. automodapi:: xrayvision.utils
diff --git a/docs/reference/visibility.rst b/docs/reference/visibility.rst
new file mode 100644
index 0000000..c3cf537
--- /dev/null
+++ b/docs/reference/visibility.rst
@@ -0,0 +1,8 @@
+.. _visibility:
+
+Visibility ('xrayvision.visibility')
+************************************
+
+The ``visibility`` submodule contains generic visibility classes.
+
+.. automodapi:: xrayvision.visibility
diff --git a/docs/tutorials/introduction.rst b/docs/tutorials/introduction.rst
index 5747c40..09ff152 100644
--- a/docs/tutorials/introduction.rst
+++ b/docs/tutorials/introduction.rst
@@ -41,10 +41,10 @@ measured visibilities in f can the original map a be recovered?
 
     data = make_data()
 
-    full_uv = transform.generate_uv(65)
+    full_uv = transform.generate_uv(65*u.pix)
 
-    uu = transform.generate_uv(33)
-    vv = transform.generate_uv(33)
+    uu = transform.generate_uv(33*u.pix)
+    vv = transform.generate_uv(33*u.pix)
 
     uu, vv = np.meshgrid(uu, vv)
 
@@ -53,7 +53,7 @@ measured visibilities in f can the original map a be recovered?
 
     full_vis = transform.dft_map(data, u=uv[0,:], v=uv[1,:])
 
-    res = transform.idft_map(full_vis, u=uv[0,:], v=uv[1,:], shape=(33, 33))
+    res = transform.idft_map(full_vis, u=uv[0,:], v=uv[1,:], weights=1/33**2, shape=(33, 33)*u.pix)
     # assert np.allclose(data, res)
 
     # Generate log spaced radial u, v sampeling
@@ -77,9 +77,10 @@ measured visibilities in f can the original map a be recovered?
     sub_vis = transform.dft_map(data, u=sub_uv[0,:], v=sub_uv[1,:])
 
     psf1 = transform.idft_map(np.full(sub_vis.size, 1), u=sub_uv[0,:], v=sub_uv[1,:],
-                              shape=(65, 65))
+                              weights=1/sub_vis.size, shape=(65, 65)*u.pix)
 
-    sub_res = transform.idft_map(sub_vis, u=sub_uv[0,:], v=sub_uv[1,:], shape=(65, 65))
+    sub_res = transform.idft_map(sub_vis, u=sub_uv[0,:], v=sub_uv[1,:],
+                                 weights=1/sub_vis.size, shape=(65, 65)*u.pix)
 
     xp = np.round(x * 33 + 33/2 - 0.5 + 16).astype(int)
     yp = np.round(y * 33 + 33/2 - 0.5 + 16).astype(int)
diff --git a/examples/rhessi.py b/examples/rhessi.py
index 79f113c..bd75884 100644
--- a/examples/rhessi.py
+++ b/examples/rhessi.py
@@ -64,13 +64,13 @@
 ###############################################################################
 # Lets have a look at the point spread function (PSF) or dirty beam
 
-psf_map = vis_psf_map(vis, shape=(101, 101) * apu.pixel, pixel_size=1.5 * apu.arcsec, natural=False)
+psf_map = vis_psf_map(vis, shape=(101, 101) * apu.pixel, pixel_size=1.5 * apu.arcsec / apu.pixel, scheme="uniform")
 
 ###############################################################################
 # We can now make an image using the back projection algorithm essentially and
 # inverse Fourier transform of the visibilities.
 
-backproj_map = vis_to_map(vis, shape=[101, 101] * apu.pixel, pixel_size=1.5 * apu.arcsec)
+backproj_map = vis_to_map(vis, shape=[101, 101] * apu.pixel, pixel_size=1.5 * apu.arcsec / apu.pix)
 
 ###############################################################################
 # Back projection contain many artifact due to the incomplete sampling of the u-v
@@ -83,13 +83,17 @@
 #                     v=vis_data_59['v']/apu.arcsec, offset=vis_data_59['xyoffset'][0]*apu.arcsec)
 
 clean_map, model_map, residual_map = vis_clean(
-    vis, shape=[101, 101] * apu.pixel, pixel=[1.5, 1.5] * apu.arcsec, clean_beam_width=10 * apu.arcsec, niter=100
+    vis,
+    shape=[101, 101] * apu.pixel,
+    pixel_size=[1.5, 1.5] * apu.arcsec / apu.pix,
+    clean_beam_width=10 * apu.arcsec,
+    niter=100,
 )
 
 ###############################################################################
 # MEM
 
-mem_map = mem(vis, shape=[129, 129] * apu.pixel, pixel=[2, 2] * apu.arcsec)
+mem_map = mem(vis, shape=[129, 129] * apu.pixel, pixel_size=[2, 2] * apu.arcsec / apu.pix)
 mem_map.plot()
 
 
diff --git a/examples/stix.py b/examples/stix.py
index 4762e67..c8511c3 100644
--- a/examples/stix.py
+++ b/examples/stix.py
@@ -26,32 +26,37 @@
 stix_data = pickle.load(urllib.request.urlopen("https://pub099.cs.technik.fhnw.ch/demo/stix_vis.pkl"))
 
 time_range, energy_range, offset, stix_vis = stix_data
+stix_vis.phase_centre = [0, 0] * apu.arcsec
 stix_vis.offset = offset
 
 ###############################################################################
 # Lets have a look at the point spread function (PSF) or dirty beam
 
-psf_map = vis_psf_map(stix_vis, shape=(129, 129) * apu.pixel, pixel_size=2 * apu.arcsec, natural=False)
+psf_map = vis_psf_map(stix_vis, shape=(129, 129) * apu.pixel, pixel_size=2 * apu.arcsec / apu.pix, scheme="uniform")
 psf_map.plot()
 
 ###############################################################################
 # Back projection
 
-backproj_map = vis_to_map(stix_vis, shape=(129, 129) * apu.pixel, pixel_size=2 * apu.arcsec, natural=False)
+backproj_map = vis_to_map(stix_vis, shape=(129, 129) * apu.pixel, pixel_size=2 * apu.arcsec / apu.pix, scheme="uniform")
 backproj_map.plot()
 
 ###############################################################################
 # Clean
 
 clean_map, model_map, resid_map = vis_clean(
-    stix_vis, shape=[129, 129] * apu.pixel, pixel=[2, 2] * apu.arcsec, niter=100, clean_beam_width=20 * apu.arcsec
+    stix_vis,
+    shape=[129, 129] * apu.pixel,
+    pixel_size=[2, 2] * apu.arcsec / apu.pix,
+    clean_beam_width=20 * apu.arcsec,
+    niter=100,
 )
 clean_map.plot()
 
 ###############################################################################
 # MEM
 
-mem_map = mem(stix_vis, shape=[129, 129] * apu.pixel, pixel=[2, 2] * apu.arcsec)
+mem_map = mem(stix_vis, shape=[129, 129] * apu.pixel, pixel_size=[2, 2] * apu.arcsec / apu.pix)
 mem_map.plot()
 
 ###############################################################################
diff --git a/setup.cfg b/setup.cfg
index a5a5c9f..28af907 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -97,6 +97,9 @@ ignore-words-list =
     process,
     technik
 
+[mypy]
+disable_error_code = import-untyped
+
 [coverage:run]
 omit =
   xrayvision/_sunpy_init*
diff --git a/xrayvision/__init__.py b/xrayvision/__init__.py
index 1623e23..612afc7 100644
--- a/xrayvision/__init__.py
+++ b/xrayvision/__init__.py
@@ -1,11 +1,7 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 
 try:
-    from xrayvision.version import __version__
+    from xrayvision.version import __version__  # type: ignore
 except ImportError:
     __version__ = "unknown"
-__all__ = []
-
-from pkg_resources import resource_filename
-
-SAMPLE_RHESSI_VISIBILITIES = resource_filename("xrayvision", "data/hsi_visibili_20131028_0156_20131028_0200_6_12.fits")
+__all__: list[str] = []
diff --git a/xrayvision/clean.py b/xrayvision/clean.py
index f10ead0..32327a3 100644
--- a/xrayvision/clean.py
+++ b/xrayvision/clean.py
@@ -8,34 +8,41 @@
 """
 
 import logging
+from typing import Union, Optional
+from collections.abc import Iterable
 
+import astropy.units as u
 import numpy as np
 from astropy.convolution import Gaussian2DKernel
+from astropy.units import Quantity
+from numpy.typing import NDArray
 from scipy import signal
-from scipy.ndimage.interpolation import shift
+from scipy.ndimage import shift
 from sunpy.map.map_factory import Map
 
 from xrayvision.imaging import vis_psf_image, vis_to_map
+from xrayvision.visibility import Visibility
 
 __all__ = ["clean", "vis_clean", "ms_clean", "vis_ms_clean"]
 
+
 logger = logging.getLogger(__name__)
 logger.setLevel("DEBUG")
 
 __common_clean_doc__ = r"""
-    clean_beam_width : `float`
+    clean_beam_width :
         The width of the gaussian to convolve the model with. If set to 0.0 \
         the gaussian to convolution is disabled
-    gain : `float`
+    gain :
         The gain per loop or loop gain
-    thres : `float`
+    thres :
         Terminates clean when ``residual.max() <= thres``
-    niter : `int`
+    niter :
         Maximum number of iterations to perform
 
     Returns
     -------
-    `numpy.ndarray`
+    :
         The CLEAN image 2D
 
     Notes
@@ -59,30 +66,39 @@
     """
 
 
-def clean(dirty_map, dirty_beam, pixel=None, clean_beam_width=4.0, gain=0.1, thres=0.01, niter=5000):
+@u.quantity_input
+def clean(
+    dirty_map: Quantity,
+    dirty_beam: Quantity,
+    pixel_size: Quantity[u.arcsec / u.pix] = None,
+    clean_beam_width: Optional[Quantity[u.arcsec]] = 4.0 * u.arcsec,
+    gain: Optional[float] = 0.1,
+    thres: Optional[float] = None,
+    niter: int = 5000,
+) -> Union[Quantity, NDArray[np.float64]]:
     r"""
     Clean the image using Hogbom's original method.
 
     CLEAN iteratively subtracts the PSF or dirty beam from the dirty map to create the residual.
-    At each iteration the location of the maximum residual is found and a shifted dirty beam is
-    subtracted that location updating the residual. This process continues until either `niter`
-    iterations is reached or the maximum residual <= `thres`.
+    At each iteration, the location of the maximum residual is found and a shifted dirty beam is
+    subtracted at that location. This process continues until either `niter` iterations are reached or
+    the maximum residual <= `thres`.
 
     Parameters
     ----------
-    dirty_map : `numpy.ndarray`
+    dirty_map :
         The dirty map to be cleaned 2D
-    dirty_beam : `numpy.ndarray`
+    dirty_beam :
         The dirty beam or point spread function (PSF) 2D must
-    pixel :
-        Size of a pixel
+    pixel_size :
+        The pixel size in arcsec
     """
     # Ensure both beam and map are even/odd on same axes
     # if not [x % 2 == 0 for x in dirty_map.shape] == [x % 2 == 0 for x in dirty_beam.shape]:
     #     raise ValueError('')
     pad = [0 if x % 2 == 0 else 1 for x in dirty_map.shape]
 
-    # Assume beam, map center is in middle
+    # Assume beam, map phase_centre is in middle
     beam_center = (dirty_beam.shape[0] - 1) / 2.0, (dirty_beam.shape[1] - 1) / 2.0
     map_center = (dirty_map.shape[0] - 1) / 2.0, (dirty_map.shape[1] - 1) / 2.0
 
@@ -104,7 +120,8 @@ def clean(dirty_map, dirty_beam, pixel=None, clean_beam_width=4.0, gain=0.1, thr
         # imax = imax * max_beam
         model[mx, my] += gain * imax
 
-        logger.debug(f"Iter: {i}, strength: {imax}, location: {mx, my}")
+        if i % 25 == 0:
+            logger.debug(f"Iter: {i}, strength: {imax}, location: {mx, my}")
 
         offset = map_center[0] - mx, map_center[1] - my
         shifted_beam_center = int(beam_center[0] + offset[0]), int(beam_center[1] + offset[1])
@@ -119,10 +136,11 @@ def clean(dirty_map, dirty_beam, pixel=None, clean_beam_width=4.0, gain=0.1, thr
 
         dirty_map = np.subtract(dirty_map, comp)
 
-        # if dirty_map.max() <= thres:
-        #      logger.info("Threshold reached")
-        #      break
-        # # el
+        if thres:
+            if np.abs(dirty_map).max() <= thres:
+                logger.info("Threshold reached")
+                break
+
         if np.abs(dirty_map.min()) > dirty_map.max():
             logger.info("Largest residual negative")
             break
@@ -130,28 +148,40 @@ def clean(dirty_map, dirty_beam, pixel=None, clean_beam_width=4.0, gain=0.1, thr
     else:
         print("Max iterations reached")
 
-    if clean_beam_width != 0.0:
+    if clean_beam_width is not None:
         # Convert from FWHM to StDev   FWHM = sigma*(8ln2)**0.5 = 2.3548200450309493
-        x_stdev = (clean_beam_width / pixel[0] / 2.3548200450309493).value
-        y_stdev = (clean_beam_width / pixel[1] / 2.3548200450309493).value
+        x_stdev = (clean_beam_width / pixel_size[1]).to_value(u.pix) / 2.3548200450309493
+        y_stdev = (clean_beam_width / pixel_size[0]).to_value(u.pix) / 2.3548200450309493
         clean_beam = Gaussian2DKernel(x_stdev, y_stdev, x_size=dirty_beam.shape[1], y_size=dirty_beam.shape[0]).array
         # Normalise beam
         clean_beam = clean_beam / clean_beam.max()
 
         # Convolve clean beam with model and scale
-        clean_map = signal.convolve2d(model, clean_beam / clean_beam.sum(), mode="same") / (pixel[0] * pixel[1])
+        clean_map = signal.convolve2d(model, clean_beam / clean_beam.sum(), mode="same") / (
+            pixel_size[0].value * pixel_size[1].value
+        )
 
         # Scale residual map with model and scale
-        dirty_map = dirty_map / clean_beam.sum() / (pixel[0] * pixel[1])
+        dirty_map = dirty_map / clean_beam.sum() / (pixel_size[0].value * pixel_size[1].value)
         return clean_map + dirty_map, model, dirty_map
 
     return model + dirty_map, model, dirty_map
 
 
-clean.__doc__ += __common_clean_doc__
+clean.__doc__ += __common_clean_doc__  # type: ignore
 
 
-def vis_clean(vis, shape, pixel, clean_beam_width=4.0, niter=5000, map=True, gain=0.1, **kwargs):
+@u.quantity_input
+def vis_clean(
+    vis: Visibility,
+    shape: Quantity[u.pix],
+    pixel_size: Quantity[u.arcsec / u.pix],
+    clean_beam_width: Optional[Quantity[u.arcsec]] = 4.0,
+    niter: Optional[int] = 5000,
+    map: Optional[bool] = True,
+    gain: Optional[float] = 0.1,
+    **kwargs,
+):
     r"""
     Clean the visibilities using Hogbom's original method.
 
@@ -159,21 +189,26 @@ def vis_clean(vis, shape, pixel, clean_beam_width=4.0, niter=5000, map=True, gai
 
     Parameters
     ----------
-    vis : `xrayvision.visibility.Visibly`
+    vis :
         The visibilities to clean
     shape :
         Size of map
-    pixel :
-        Size of pixel
-    map : `boolean` optional
-        Return an `sunpy.map.Map` by default or data only if `False`
+    pixel_size :
+        The pixel size in arcsec
+    map :
+        Return a `sunpy.map.Map` by default or array only if `False`
     """
 
-    dirty_map = vis_to_map(vis, shape=shape, pixel_size=pixel, **kwargs)
+    dirty_map = vis_to_map(vis, shape=shape, pixel_size=pixel_size, **kwargs)
     dirty_beam_shape = [x.value * 3 + 1 if x.value * 3 % 2 == 0 else x.value * 3 for x in shape] * shape.unit
-    dirty_beam = vis_psf_image(vis, shape=dirty_beam_shape, pixel_size=pixel, **kwargs)
+    dirty_beam = vis_psf_image(vis, shape=dirty_beam_shape, pixel_size=pixel_size, **kwargs)
     clean_map, model, residual = clean(
-        dirty_map.data, dirty_beam.value, pixel=pixel, gain=gain, clean_beam_width=clean_beam_width, niter=niter
+        dirty_map.data,
+        dirty_beam.value,
+        pixel_size=pixel_size,
+        clean_beam_width=clean_beam_width,
+        gain=gain,
+        niter=niter,
     )
     if not map:
         return clean_map, model, residual
@@ -181,25 +216,24 @@ def vis_clean(vis, shape, pixel, clean_beam_width=4.0, niter=5000, map=True, gai
     return [Map((data, dirty_map.meta)) for data in (clean_map, model, residual)]
 
 
-vis_clean.__doc__ += __common_clean_doc__
-
+vis_clean.__doc__ += __common_clean_doc__  # type: ignore
 
 __common_ms_clean_doc__ = r"""
     scales : array-like, optional, optional
         The scales to use eg ``[1, 2, 4, 8]``
-    clean_beam_width : `float`
+    clean_beam_width :
         The width of the gaussian to convolve the model with. If set to 0.0 the gaussian \
         convolution is disabled
-    gain : `float`
+    gain :
         The gain per loop or loop gain
-    thres : `float`
+    thres :
         Terminates clean when `residuals.max() <= thres``
-    niter : `int`
+    niter :
         Maximum number of iterations to perform
 
     Returns
     -------
-    `numpy.ndarray`
+    :
         Cleaned image
 
     Notes
@@ -219,20 +253,32 @@ def vis_clean(vis, shape, pixel, clean_beam_width=4.0, niter=5000, map=True, gai
     """
 
 
-def ms_clean(dirty_map, dirty_beam, pixel, scales=None, clean_beam_width=4.0, gain=0.1, thres=0.01, niter=5000):
+@u.quantity_input
+def ms_clean(
+    dirty_map: Quantity,
+    dirty_beam: Quantity,
+    pixel_size: Quantity[u.arcsec / u.pix],
+    scales: Union[Iterable, NDArray, None] = None,
+    clean_beam_width: Quantity = 4.0 * u.arcsec,
+    gain: float = 0.1,
+    thres: float = 0.01,
+    niter: int = 5000,
+) -> Union[Quantity, NDArray[np.float64]]:
     r"""
     Clean the map using a multiscale clean algorithm.
 
     Parameters
     ----------
-    dirty_map : `numpy.ndarray`
+    dirty_map :
         The 2D dirty map to be cleaned
-    dirty_beam : `numpy.ndarray`
+    dirty_beam :
         The 2D dirty beam should have the same dimensions as `dirty_map`
+    pixel_size :
+        The pixel size in arcsec
     """
     # Compute the number of dyadic scales, their sizes and scale biases
-    number_of_scales = np.floor(np.log2(min(dirty_map.shape))).astype(int)
-    scale_sizes = 2 ** np.arange(number_of_scales)
+    number_of_scales: int = np.floor(np.log2(min(dirty_map.shape))).astype(int)
+    scale_sizes: NDArray[np.int_] = 2 ** np.arange(number_of_scales)
 
     if scales:
         scales = np.array(scales)
@@ -258,7 +304,7 @@ def ms_clean(dirty_map, dirty_beam, pixel, scales=None, clean_beam_width=4.0, ga
     cross_terms = {}
 
     for i, scale in enumerate(scale_sizes):
-        scales[:, :, i] = component(scale=scale, shape=dirty_map.shape)
+        scales[:, :, i] = _component(scale=scale, shape=dirty_map.shape)
         scaled_residuals[:, :, i] = signal.convolve(dirty_map, scales[:, :, i], mode="same")
         scaled_dirty_beams[:, :, i] = signal.convolve(dirty_beam, scales[:, :, i], mode="same")
         max_scaled_dirty_beams[i] = scaled_dirty_beams[:, :, i].max()
@@ -272,8 +318,11 @@ def ms_clean(dirty_map, dirty_beam, pixel, scales=None, clean_beam_width=4.0, ga
         # print(f'Clean loop {i}')
         # For each scale find the strength and location of max residual
         # Chose scale with has maximum strength
-        max_index = np.argmax(scaled_residuals)
-        max_x, max_y, max_scale = np.unravel_index(max_index, scaled_residuals.shape)
+        max_index: int = np.argmax(scaled_residuals).astype(int)
+        max_x: int
+        max_y: int
+        max_scale: int
+        max_x, max_y, max_scale = (int(x) for x in np.unravel_index(max_index, scaled_residuals.shape))
 
         strength = scaled_residuals[max_x, max_y, max_scale]
 
@@ -319,9 +368,9 @@ def ms_clean(dirty_map, dirty_beam, pixel, scales=None, clean_beam_width=4.0, ga
             scaled_residuals[:, :, j] = np.subtract(scaled_residuals[:, :, j], comp)
 
         # End max(res(a)) or niter
-        if scaled_residuals[:, :, max_scale].max() <= thres:
+        if np.abs(scaled_residuals[:, :, max_scale].max()) <= thres:
             logger.info("Threshold reached")
-            # break
+            break
 
         # Largest scales largest residual is negative
         if np.abs(scaled_residuals[:, :, 0].min()) > scaled_residuals[:, :, 0].max():
@@ -331,29 +380,39 @@ def ms_clean(dirty_map, dirty_beam, pixel, scales=None, clean_beam_width=4.0, ga
     else:
         logger.info("Max iterations reached")
 
-    # Convolve model with clean beam  B_G * I^M
-    if clean_beam_width != 0.0:
-        x_stdev = (clean_beam_width / pixel[0] / (2.0 * np.sqrt(2.0 * np.log(2.0)))).value
-        y_stdev = (clean_beam_width / pixel[1] / (2.0 * np.sqrt(2.0 * np.log(2.0)))).value
+    # Convolve model with clean beam B_G * I^M
+    if clean_beam_width is not None:
+        x_stdev = ((clean_beam_width / pixel_size[0]).to_value(u.pix) / (2.0 * np.sqrt(2.0 * np.log(2.0)))).value
+        y_stdev = ((clean_beam_width / pixel_size[1]).to_value(u.pix) / (2.0 * np.sqrt(2.0 * np.log(2.0)))).value
         clean_beam = Gaussian2DKernel(x_stdev, y_stdev, x_size=dirty_beam.shape[1], y_size=dirty_beam.shape[0]).array
 
         # Normalise beam
         clean_beam = clean_beam / clean_beam.max()
 
-        clean_map = signal.convolve2d(model, clean_beam, mode="same") / (pixel[0] * pixel[1])
+        clean_map = signal.convolve2d(model, clean_beam, mode="same") / (pixel_size[0] * pixel_size[1])
 
         # Scale residual map with model and scale
-        dirty_map = (scaled_residuals / clean_beam.sum() / (pixel[0] * pixel[1])).sum(axis=2)
+        dirty_map = (scaled_residuals / clean_beam.sum() / (pixel_size[0] * pixel_size[1])).sum(axis=2)
 
         return clean_map + dirty_map, model, dirty_map
     # Add residuals B_G * I^M + I^R
     return model, scaled_residuals.sum(axis=2)
 
 
-ms_clean.__doc__ += __common_ms_clean_doc__
+ms_clean.__doc__ += __common_ms_clean_doc__  # type: ignore
 
 
-def vis_ms_clean(vis, shape, pixel, scales=None, clean_beam_width=4.0, gain=0.1, thres=0.01, niter=5000, map=True):
+def vis_ms_clean(
+    vis: Visibility,
+    shape: Quantity[u.pix],
+    pixel_size: Quantity[u.arcsec / u.pix],
+    scales: Optional[Iterable],
+    clean_beam_width: Optional[Quantity[u.arcsec]] = 4.0,
+    niter: Optional[int] = 5000,
+    map: Optional[bool] = True,
+    gain: Optional[float] = 0.1,
+    thres: Optional[float] = 0.01,
+) -> Union[Quantity, NDArray[np.float64]]:
     r"""
     Clean the visibilities using a multiscale clean method.
 
@@ -361,20 +420,38 @@ def vis_ms_clean(vis, shape, pixel, scales=None, clean_beam_width=4.0, gain=0.1,
 
     Parameters
     ----------
-    vis : `xrayvision.visibility.Visibly`
+    vis :
         The visibilities to clean
     shape :
         Size of map
-    pixel :
-        Size of pixel
+    pixel_size :
+        The pixel size in arcsec
+    scales : array-like, optional, optional
+        The scales to use eg ``[1, 2, 4, 8]``
+    clean_beam_width :
+        The width of the gaussian to convolve the model with. If set to 0.0 the gaussian \
+        convolution is disabled
+    gain :
+        The gain per loop or loop gain
+    thres :
+        Terminates clean when `residuals.max() <= thres``
+    niter :
+        Maximum number of iterations to perform
+    map :
+        Return a `sunpy.map.Map` by default or array only if `False`
+
+    Returns
+    -------
+    :
+        Cleaned image
 
     """
-    dirty_map = vis_to_map(vis, shape=shape, pixel_size=pixel)
-    dirty_beam = vis_psf_image(vis, shape=shape * 3, pixel_size=pixel)
+    dirty_map = vis_to_map(vis, shape=shape, pixel_size=pixel_size)
+    dirty_beam = vis_psf_image(vis, shape=shape * 3, pixel_size=pixel_size)
     clean_map, model, residual = ms_clean(
         dirty_map.data,
         dirty_beam,
-        pixel,
+        pixel_size=pixel_size,
         scales=scales,
         clean_beam_width=clean_beam_width,
         gain=gain,
@@ -390,7 +467,7 @@ def vis_ms_clean(vis, shape, pixel, scales=None, clean_beam_width=4.0, gain=0.1,
 # vis_ms_clean.__doc__ += __common_ms_clean_doc__
 
 
-def radial_prolate_sphereoidal(nu):
+def _radial_prolate_sphereoidal(nu):
     r"""
     Calculate prolate spheroidal wave function approximation.
 
@@ -467,7 +544,7 @@ def radial_prolate_sphereoidal(nu):
             return 0
 
 
-def vec_radial_prolate_sphereoidal(nu):
+def _vec_radial_prolate_sphereoidal(nu):
     r"""
     Calculate prolate spheroidal wave function approximation.
 
@@ -546,7 +623,7 @@ def vec_radial_prolate_sphereoidal(nu):
     return out
 
 
-def component(scale, shape):
+def _component(scale, shape):
     r"""
 
     Parameters
@@ -579,9 +656,9 @@ def component(scale, shape):
     rad_zeros_indices = radii_squared <= 0.0
     amp_zero_indices = radii_squared >= 1.0
 
-    wave_amp = vec_radial_prolate_sphereoidal(np.sqrt(radii_squared.reshape(radii_squared.size)))
+    wave_amp = _vec_radial_prolate_sphereoidal(np.sqrt(radii_squared.reshape(radii_squared.size)))
     wave_amp = wave_amp.reshape(shape)
-    wave_amp[rad_zeros_indices] = vec_radial_prolate_sphereoidal([0])[0]
+    wave_amp[rad_zeros_indices] = _vec_radial_prolate_sphereoidal([0])[0]
 
     wave_amp = wave_amp * (1 - radii_squared)
 
diff --git a/xrayvision/imaging.py b/xrayvision/imaging.py
index d1559d7..c304266 100644
--- a/xrayvision/imaging.py
+++ b/xrayvision/imaging.py
@@ -1,6 +1,9 @@
+from typing import Optional
+
 import astropy.units as apu
 import numpy as np
-from sunpy.map import Map
+from astropy.units import Quantity
+from sunpy.map import GenericMap, Map
 
 from xrayvision.transform import dft_map, idft_map
 from xrayvision.visibility import Visibility
@@ -18,30 +21,32 @@
 ]
 
 ANGLE = apu.get_physical_type(apu.deg)
+WEIGHT_SCHEMES = ("natural", "uniform")
 
 
-def get_weights(vis, natural=True, norm=True):
+def get_weights(vis: Visibility, scheme: str = "natural", norm: bool = True) -> np.ndarray:
     r"""
-    Return natural spatial frequency weight factor for each visibility.
+    Return spatial frequency weight factors for each visibility.
 
     Defaults to use natural weighting scheme given by `(vis.u**2 + vis.v**2)^{1/2}`.
 
     Parameters
     ----------
-    vis : `xrayvision.visibility.Visibility`
+    vis :
         Input visibilities
-    natural : `boolean` optional
-        Use natural or uniform weighting scheme
-    norm : `boolean`
-        Normalise the weighs before returning
+    scheme :
+        Weighting scheme to use, defaults to natural
+    norm :
+        Normalise the weighs before returning, defaults to True.
 
     Returns
     -------
-    `weights`
 
     """
+    if scheme not in WEIGHT_SCHEMES:
+        raise ValueError(f"Invalid weighting scheme {scheme}, must be one of: {WEIGHT_SCHEMES}")
     weights = np.sqrt(vis.u**2 + vis.v**2).value
-    if natural:
+    if scheme == "natural":
         weights = np.ones_like(vis.vis, dtype=float)
 
     if norm:
@@ -50,7 +55,8 @@ def get_weights(vis, natural=True, norm=True):
     return weights
 
 
-def validate_and_expand_kwarg(q, name=""):
+@apu.quantity_input()
+def validate_and_expand_kwarg(q: Quantity, name: Optional[str] = "") -> Quantity:
     r"""
     Expand a scalar or array of size one to size two by repeating.
 
@@ -85,223 +91,243 @@ def validate_and_expand_kwarg(q, name=""):
     return q
 
 
-@apu.quantity_input(shape=apu.pixel, pixel_size="angle")
-def vis_psf_image(vis, *, shape=(65, 65) * apu.pixel, pixel_size=2 * apu.arcsec, natural=True):
+@apu.quantity_input
+def image_to_vis(
+    image: Quantity,
+    *,
+    u: Quantity[apu.arcsec**-1],
+    v: Quantity[apu.arcsec**-1],
+    phase_centre: Optional[Quantity[apu.arcsec]] = (0.0, 0.0) * apu.arcsec,
+    pixel_size: Optional[Quantity[apu.arcsec / apu.pix]] = 1.0 * apu.arcsec / apu.pix,
+) -> Visibility:
+    r"""
+    Return a Visibility created from the image and u, v sampling.
+
+    Parameters
+    ----------
+    image :
+        The 2D input image
+    u :
+        Array of u coordinates where the visibilities will be evaluated
+    v :
+        Array of v coordinates where the visibilities will be evaluated
+    phase_centre :
+        The coordinates the phase_centre.
+    pixel_size :
+        Size of pixels, if only one value is passed, assume square pixels (repeating the value).
+
+    Returns
+    -------
+    :
+        The new visibility object
+
     """
-    Create the point spread function for given u, v point of the visibilities.
+    pixel_size = validate_and_expand_kwarg(pixel_size, "pixel_size")
+    if not (apu.get_physical_type((1 / u).unit) == ANGLE and apu.get_physical_type((1 / v).unit) == ANGLE):
+        raise ValueError("u and v must be inverse angle (e.g. 1/deg or 1/arcsec")
+    vis = dft_map(image, u=u, v=v, phase_centre=phase_centre, pixel_size=pixel_size)
+    return Visibility(vis, u=u, v=v, offset=phase_centre)
+
+
+@apu.quantity_input()
+def vis_to_image(
+    vis: Visibility,
+    shape: Quantity[apu.pix] = (65, 65) * apu.pixel,
+    pixel_size: Optional[Quantity[apu.arcsec / apu.pix]] = 1 * apu.arcsec / apu.pix,
+    scheme: str = "natural",
+) -> Quantity:
+    """
+    Create an image by 'back projecting' the given visibilities onto the sky.
 
     Parameters
     ----------
-    vis : `xrayvision.visibility.Visibility`
+    vis :
         Input visibilities
-    shape : `astropy.units.quantity.Quantity`, optional
-        Shape of the image, if only one value is given assume square (repeating the value).
-    pixel_size : `astropy.units.quantity.Quantity`, optional
-        Size of pixels, if only one value is given assume square pixels (repeating the value).
-    natural : `boolean` optional
-        Weight scheme use natural by default, uniform if `False`
+    shape :
+        Shape of the image, if only one value is passed assume square (repeating the value).
+    pixel_size :
+        Size of pixels, if only one value is passed assume square pixels (repeating the value).
+    scheme :
+        Weight scheme natural by default.
 
     Returns
     -------
-    `astropy.units.quantity.Quantity`
-        Point spread function
+    `~astropy.units.Quantity`
+        Back projection image
 
     """
     shape = validate_and_expand_kwarg(shape, "shape")
     pixel_size = validate_and_expand_kwarg(pixel_size, "pixel_size")
-    shape = shape.to_value(apu.pixel)
-    weights = get_weights(vis, natural=natural)
-
-    # Make sure psf is always odd so power is in exactly one pixel
-    m, n = (s // 2 * 2 + 1 for s in shape)
-    psf_arr = idft_map(
-        np.ones(vis.vis.shape) * vis.vis.unit, u=vis.u, v=vis.v, shape=(m, n), weights=weights, pixel_size=pixel_size
+    shape = shape.to(apu.pixel)
+    weights = get_weights(vis, scheme=scheme)
+    bp_arr = idft_map(
+        vis.vis, u=vis.u, v=vis.v, shape=shape, weights=weights, pixel_size=pixel_size, phase_centre=vis.phase_centre
     )
-    return psf_arr
+
+    return bp_arr
 
 
-@apu.quantity_input(shape=apu.pixel, pixel_size="angle")
-def vis_psf_map(vis, *, shape=(65, 65) * apu.pixel, pixel_size=2 * apu.arcsec, natural=True):
+@apu.quantity_input
+def vis_psf_map(
+    vis: Visibility,
+    *,
+    shape: Quantity[apu.pix] = (65, 65) * apu.pixel,
+    pixel_size: Optional[Quantity[apu.arcsec / apu.pix]] = 1 * apu.arcsec / apu.pix,
+    scheme: Optional[str] = "natural",
+) -> GenericMap:
     r"""
     Create a map of the point spread function for given the visibilities.
 
     Parameters
     ----------
-    vis : `xrayvision.visibility.Visibility`
+    vis :
         Input visibilities
-    shape : `astropy.units.quantity.Quantity`, optional
-        Shape of the image, if only one value is given assume square (repeating the value).
-    pixel_size : `Tuple[~astropy.units.Quantity]`, optional
-        Size of pixels, if only one value is given assume square pixels (repeating the value).
-    natural : `boolean` optional
-        Weight scheme use natural by default, uniform if `False`
+    shape :
+        Shape of the image, if only one value is passed, assume square (repeating the value).
+    pixel_size :
+        Size of pixels, if only one value is passed, assume square pixels (repeating the value).
+    scheme :
+        Weight scheme to use natural by default.
 
     Returns
     -------
-
+    :
+        Map of the point spread function
     """
     shape = validate_and_expand_kwarg(shape, "shape")
     pixel_size = validate_and_expand_kwarg(pixel_size, "pixel_size")
     header = generate_header(vis, shape=shape, pixel_size=pixel_size)
-    psf = vis_psf_image(vis, shape=shape, pixel_size=pixel_size, natural=natural)
+    psf = vis_psf_image(vis, shape=shape, pixel_size=pixel_size, scheme=scheme)
     return Map((psf, header))
 
 
-@apu.quantity_input(shape=apu.pixel, pixel_size="angle")
-def vis_to_image(vis, shape=(65, 65) * apu.pixel, pixel_size=2 * apu.arcsec, natural=True):
+@apu.quantity_input()
+def vis_psf_image(
+    vis: Visibility,
+    *,
+    shape: Quantity[apu.pix] = (65, 65) * apu.pixel,
+    pixel_size: Quantity[apu.arcsec / apu.pix] = 1 * apu.arcsec / apu.pix,
+    scheme: str = "natural",
+) -> Quantity:
     """
-    Create an image by 'back projecting' the given visibilities onto the sky.
+    Create the point spread function for given u, v point of the visibilities.
 
     Parameters
     ----------
-    vis : `xrayvision.visibility.Visibility`
+    vis :
         Input visibilities
-    shape : `~astropy.units.Quantity`
-        Shape of the image, if only one value is given assume square (repeating the value).
-    pixel_size : `~astropy.units.Quantity`
-        Size of pixels, if only one value is given assume square pixels (repeating the value).
-    natural : `boolean` optional
-        Weight scheme use natural by default, uniform if `False`
+    shape :
+        Shape of the image, if only one value passed, assume square (repeating the value).
+    pixel_size :
+        Size of pixels, if only one value is passed, assume square pixels (repeating the value).
+    scheme :
+        Weight scheme, natural by default.
 
     Returns
     -------
-    `~astropy.units.Quantity`
-        Back projection image
+    :
+        Point spread function
 
     """
     shape = validate_and_expand_kwarg(shape, "shape")
     pixel_size = validate_and_expand_kwarg(pixel_size, "pixel_size")
-    shape = shape.to_value(apu.pixel)
-    weights = get_weights(vis, natural=natural)
-    bp_arr = idft_map(vis.vis, u=vis.u, v=vis.v, shape=shape, weights=weights, pixel_size=pixel_size, center=vis.center)
+    shape = shape.to(apu.pixel)
+    weights = get_weights(vis, scheme=scheme)
 
-    return bp_arr
+    # Make sure psf is always odd so power is in exactly one pixel
+    shape = [s // 2 * 2 + 1 for s in shape.to_value(apu.pix)] * shape.unit
+    psf_arr = idft_map(
+        np.ones(vis.vis.shape) * vis.vis.unit, u=vis.u, v=vis.v, shape=shape, weights=weights, pixel_size=pixel_size
+    )
+    return psf_arr
 
 
-@apu.quantity_input(shape=apu.pixel, pixel_size="angle")
-def vis_to_map(vis, shape=(65, 65) * apu.pixel, pixel_size=2 * apu.arcsec, natural=True):
+@apu.quantity_input()
+def vis_to_map(
+    vis: Visibility,
+    shape: Quantity[apu.pix] = (65, 65) * apu.pixel,
+    pixel_size: Optional[Quantity[apu.arcsec / apu.pix]] = 1 * apu.arcsec / apu.pixel,
+    scheme: Optional[str] = "natural",
+) -> GenericMap:
     r"""
     Create a map by performing a back projection of inverse transform on the visibilities.
 
     Parameters
     ----------
-    vis : `xrayvision.visibility.Visibility`
+    vis :
         Input visibilities
-    shape : `int` (m, n)
-        Shape of the image, if only one value is given assume square (repeating the value).
-    pixel_size : `float` (dx, dy), optional
-        Size of pixels, if only one value is given assume square pixels (repeating the value).
-    natural : `boolean` optional
-        Weight scheme use natural by default, uniform if `False`.
+    shape :
+        Shape of the image, if only one value is passed, assume square (repeating the value).
+    pixel_size :
+        Size of pixels, if only one value is passed, assume square pixels (repeating the value).
+    scheme :
+        Weight scheme natural by default.
 
     Returns
     -------
-    `sunpy.map.Map`
-        Map object with the map created from the visibilities and the metadata will contain the
-        offset and the pixel size
+    :
+        Map object with the map created from the visibilities and the metadata will contain the offset and the pixel size
 
     """
     shape = validate_and_expand_kwarg(shape, "shape")
     pixel_size = validate_and_expand_kwarg(pixel_size, "pixel_size")
     header = generate_header(vis, shape=shape, pixel_size=pixel_size)
 
-    image = vis_to_image(vis, shape=shape, pixel_size=pixel_size, natural=natural)
+    image = vis_to_image(vis, shape=shape, pixel_size=pixel_size, scheme=scheme)
     return Map((image, header))
 
 
-def generate_header(vis, *, shape, pixel_size):
+@apu.quantity_input()
+def generate_header(vis: Visibility, *, shape: Quantity[apu.pix], pixel_size: Quantity[apu.arcsec / apu.pix]) -> dict:
     r"""
     Generate a map head given the visibilities, pixel size and shape
 
     Parameters
     ----------
     vis :
-    shape : `~astropy.units.Quantity`
-        Shape of the image, if only one value is given assume square (repeating the value).
-    pixel_size : `~astropy.units.Quantity`
-        Size of pixels, if only one value is given assume square pixels (repeating the value)
-    offset
+        Input visibilities.
+    shape :
+        Shape of the image, if only one value is passed assume square (repeating the value).
+    pixel_size :
+        Size of pixels, if only one value is passed assume square pixels (repeating the value)
 
     Returns
     -------
-
+    :
     """
     header = {
-        "crval1": (vis.offset[0]).to_value(apu.arcsec),
-        "crval2": (vis.offset[1]).to_value(apu.arcsec),
-        "cdelt1": pixel_size[0].to_value(apu.arcsec),
-        "cdelt2": pixel_size[1].to_value(apu.arcsec),
+        "crval1": (vis.offset[1]).to_value(apu.arcsec),
+        "crval2": (vis.offset[0]).to_value(apu.arcsec),
+        "cdelt1": (pixel_size[1] * apu.pix).to_value(apu.arcsec),
+        "cdelt2": (pixel_size[0] * apu.pix).to_value(apu.arcsec),
         "ctype1": "HPLN-TAN",
         "ctype2": "HPLT-TAN",
         "naxis": 2,
-        "naxis1": shape[0].value,
-        "naxis2": shape[1].value,
+        "naxis1": shape[1].value,
+        "naxis2": shape[0].value,
         "cunit1": "arcsec",
         "cunit2": "arcsec",
     }
-    # if vis.center is not None:
-    #     header['crval1'] = vis.center[0].value
-    #     header['crval2'] = vis.center[1].value
-    if pixel_size is not None:
-        if pixel_size.ndim == 0:
-            header["cdelt1"] = pixel_size.value
-            header["cdelt2"] = pixel_size.value
-        elif pixel_size.ndim == 1 and pixel_size.size == 2:
-            header["cdelt1"] = pixel_size[0].value
-            header["cdelt2"] = pixel_size[1].value
-        else:
-            raise ValueError(f"pixel_size can have a length of 1 or 2 not {pixel_size.shape}")
     return header
 
 
-@apu.quantity_input(center="angle", pixel_size="angle")
-def image_to_vis(image, *, u, v, center=(0.0, 0.0) * apu.arcsec, pixel_size=2.0 * apu.arcsec):
-    r"""
-    Return a Visibility created from the image and u, v sampling.
-
-    Parameters
-    ----------
-    image : `numpy.ndarray`
-        The 2D input image
-    u : `astropy.units.Quantity`
-        Array of u coordinates where the visibilities will be evaluated
-    v : `astropy.units.Quantity`
-        Array of v coordinates where the visibilities will be evaluated
-    center : `~astropy.units.Quantity` (x, y)
-        The coordinates of the center of the image.
-    pixel_size : `~astropy.units.Quantity`
-        Size of pixels, if only one value is given assume square pixels (repeating the value).
-
-    Returns
-    -------
-    `Visibility`
-
-        The new visibility object
-
-    """
-    pixel_size = validate_and_expand_kwarg(pixel_size, "pixel_size")
-    if not (apu.get_physical_type((1 / u).unit) == ANGLE and apu.get_physical_type((1 / v).unit) == ANGLE):
-        raise ValueError("u and v must be inverse angle (e.g. 1/deg or 1/arcsec")
-    vis = dft_map(image, u=u, v=v, center=center, pixel_size=pixel_size)
-    return Visibility(vis, u=u, v=v, center=center)
-
-
-def map_to_vis(amap, *, u, v):
+@apu.quantity_input()
+def map_to_vis(amap: GenericMap, *, u: Quantity[1 / apu.arcsec], v: Quantity[1 / apu.arcsec]) -> Visibility:
     r"""
-    Return a Visibility object created from the map and u, v sampling.
+    Return a Visibility object created from the map, sampling it at give `u`, `v` coordinates.
 
     Parameters
     ----------
-    amap : `sunpy.map.Map`
-        The input map
-    u : `numpy.ndarray`
+    amap :
+        Input map.
+    u :
         Array of u coordinates where the visibilities will be evaluated
-    v : `numpy.ndarray`
+    v :
         Array of v coordinates where the visibilities will be evaluated
 
     Returns
     -------
-    `Visibility`
+    :
         The new visibility object
 
     """
@@ -310,16 +336,16 @@ def map_to_vis(amap, *, u, v):
     meta = amap.meta
     new_pos = np.array([0.0, 0.0])
     if "crval1" in meta:
-        new_pos[0] = float(meta["crval1"])
+        new_pos[1] = float(meta["crval1"])
     if "crval2" in meta:
-        new_pos[1] = float(meta["crval2"])
+        new_pos[0] = float(meta["crval2"])
 
     new_psize = np.array([1.0, 1.0])
     if "cdelt1" in meta:
-        new_psize[0] = float(meta["cdelt1"])
+        new_psize[1] = float(meta["cdelt1"])
     if "cdelt2" in meta:
-        new_psize[1] = float(meta["cdelt2"])
+        new_psize[0] = float(meta["cdelt2"])
 
-    vis = image_to_vis(amap.data, u=u, v=v, pixel_size=new_psize * apu.arcsec)
+    vis = image_to_vis(amap.data, u=u, v=v, pixel_size=new_psize * apu.arcsec / apu.pix)
     vis.offset = new_pos * apu.arcsec
     return vis
diff --git a/xrayvision/mem.py b/xrayvision/mem.py
index 1ee4916..485e86f 100644
--- a/xrayvision/mem.py
+++ b/xrayvision/mem.py
@@ -3,10 +3,13 @@
 """
 
 from types import SimpleNamespace
+from typing import Union, Optional
 
 import astropy.units as apu
 import numpy as np
+from astropy.units import Quantity
 from numpy.linalg import norm
+from numpy.typing import NDArray
 from sunpy.map import Map
 
 from xrayvision.imaging import generate_header
@@ -14,21 +17,22 @@
 from xrayvision.utils import get_logger
 
 __all__ = [
-    "get_entropy",
-    "get_fourier_matrix",
-    "estimate_flux",
-    "get_mean_visibilities",
-    "proximal_entropy",
-    "proximal_operator",
-    "optimise_fb",
+    "_get_entropy",
+    "_get_fourier_matrix",
+    "_estimate_flux",
+    "_get_mean_visibilities",
+    "_proximal_entropy",
+    "_proximal_operator",
+    "_optimise_fb",
     "mem",
 ]
 
+from xrayvision.visibility import Visibility
 
 logger = get_logger(__name__, "DEBUG")
 
 
-def get_entropy(image, flux):
+def _get_entropy(image, flux):
     r"""
     Return the entropy of an image.
 
@@ -55,7 +59,7 @@ def get_entropy(image, flux):
     return np.sum(image * np.log(image / (flux * np.e)))
 
 
-def get_fourier_matrix(vis, shape=[64, 64] * apu.pix, pixel_size=[4.0312500, 4.0312500] * apu.arcsec):
+def _get_fourier_matrix(vis, shape=(64, 64) * apu.pix, pixel_size=(4.0312500, 4.0312500) * apu.arcsec):
     r"""
     Return the complex Fourier matrix used to compute the value of the visibilities.
 
@@ -76,10 +80,10 @@ def get_fourier_matrix(vis, shape=[64, 64] * apu.pix, pixel_size=[4.0312500, 4.0
     -------
     The complex Fourier matrix
     """
-    m, n = shape.to_value("pix")
-    y = generate_xy(m, 0 * apu.arcsec, pixel_size[1])
-    x = generate_xy(n, 0 * apu.arcsec, pixel_size[0])
-    x, y = np.meshgrid(x, y)
+    m, n = shape.to("pix")
+    y = generate_xy(m, phase_centre=0 * apu.arcsec, pixel_size=pixel_size[1])
+    x = generate_xy(n, phase_centre=0 * apu.arcsec, pixel_size=pixel_size[0])
+    x, y = np.meshgrid(x, y, indexing="ij")
     uv = np.vstack([vis.u, vis.v])
     # Check apu are correct for exp need to be dimensionless and then remove apu for speed
     if (vis.u * x[0, 0]).unit == apu.dimensionless_unscaled and (vis.v * y[0, 0]).unit == apu.dimensionless_unscaled:
@@ -91,10 +95,10 @@ def get_fourier_matrix(vis, shape=[64, 64] * apu.pix, pixel_size=[4.0312500, 4.0
             1j * 2 * np.pi * (x[..., np.newaxis] * uv[np.newaxis, 0, :] + y[..., np.newaxis] * uv[np.newaxis, 1, :])
         )
 
-        return Hv * pixel_size[0] * pixel_size[1]
+        return Hv * pixel_size[0].value * pixel_size[1].value
 
 
-def estimate_flux(vis, shape, pixel, maxiter=1000, tol=1e-3):
+def _estimate_flux(vis, shape, pixel, maxiter=1000, tol=1e-3):
     r"""
     Estimate the total flux in the image by solving an optimisation problem.
 
@@ -151,7 +155,9 @@ def estimate_flux(vis, shape, pixel, maxiter=1000, tol=1e-3):
         diff_V = Hvx - Visib
         chi2 = (diff_V**2.0).sum()
 
-        logger.info(f"Iter: {i}, Chi2: {chi2}")
+        if i % 25 == 0:
+            logger.info(f"Iter: {i}, Chi2: {chi2}")
+
         if np.sqrt(((x - x_old) ** 2.0).sum()) < tol * np.sqrt((x_old**2.0).sum()):
             break
 
@@ -179,7 +185,7 @@ def _prepare_for_optimise(pixel, shape, vis):
     -------
 
     """
-    Hv = get_fourier_matrix(vis, shape, pixel)
+    Hv = _get_fourier_matrix(vis, shape, pixel)
     # Division of real and imaginary part of the matrix 'Hv'
     ReHv = np.real(Hv)
     ImHv = np.imag(Hv)
@@ -210,7 +216,7 @@ def _prepare_for_optimise(pixel, shape, vis):
     return Hv, Lip, Visib
 
 
-def get_mean_visibilities(vis, shape, pixel):
+def _get_mean_visibilities(vis, shape, pixel):
     r"""
     Return the mean visibilities sampling the same call in the discretisation of the (u,v) plane.
 
@@ -237,9 +243,9 @@ def get_mean_visibilities(vis, shape, pixel):
     rv = np.around(iv)
 
     # index of the u coordinates of the sampling frequencies in the discretisation of the u axis
-    ru = ru + imsize2
+    ru = ru * apu.pix + imsize2.to(apu.pix)
     # index of the v coordinates of the sampling frequencies in the discretisation of the v axis
-    rv = rv + imsize2
+    rv = rv * apu.pix + imsize2.to(apu.pix)
 
     # matrix that represents the discretization of the (u,v)-plane
     iuarr = np.zeros(shape.to_value("pixel").astype(int))
@@ -293,7 +299,7 @@ def get_mean_visibilities(vis, shape, pixel):
     return SimpleNamespace(u=u, v=v, vis=visib, amplitude_error=weights)
 
 
-def proximal_entropy(y, m, lamba, Lip, tol=10**-10):
+def _proximal_entropy(y, m, lamba, Lip, tol=10**-10):
     r"""
     This function computes the value of the proximity operator of the entropy function subject to
     positivity constraint, i.e. it solves the problem
@@ -335,7 +341,7 @@ def proximal_entropy(y, m, lamba, Lip, tol=10**-10):
     return c
 
 
-def proximal_operator(z, f, m, lamb, Lip, niter=250):
+def _proximal_operator(z, f, m, lamb, Lip, niter=250):
     r"""
     Computes the value of the proximity operator of the entropy function subject to
     positivity constraint and flux constraint by means of a Dykstra-like proximal algorithm
@@ -371,7 +377,7 @@ def proximal_operator(z, f, m, lamb, Lip, niter=250):
         y = tmp + (f - tmp.sum()) / tmp.size
         p = x + p - y
 
-        x = proximal_entropy(y + q, m, lamb, Lip)
+        x = _proximal_entropy(y + q, m, lamb, Lip)
 
         if np.abs(x.sum() - f) <= 0.01 * f:
             break
@@ -381,7 +387,7 @@ def proximal_operator(z, f, m, lamb, Lip, niter=250):
     return x, i
 
 
-def optimise_fb(Hv, Visib, Lip, flux, lambd, shape, pixel, maxiter, tol):
+def _optimise_fb(Hv, Visib, Lip, flux, lambd, shape, pixel, maxiter, tol):
     r"""
     Solve the optimization problem using a forward-backward splitting algorithm
 
@@ -441,7 +447,7 @@ def optimise_fb(Hv, Visib, Lip, flux, lambd, shape, pixel, maxiter, tol):
 
     tmp = x.flatten()[:]
     Hvx = np.matmul(Hv, tmp)
-    f_R = get_entropy(x, m)
+    f_R = _get_entropy(x, m)
 
     diff_V = Hvx - Visib
     f_0 = (diff_V**2).sum()
@@ -458,12 +464,12 @@ def optimise_fb(Hv, Visib, Lip, flux, lambd, shape, pixel, maxiter, tol):
         y = z - 1 / Lip * grad
 
         # PROXIMAL STEP
-        p, pi = proximal_operator(y, f, m, lambd, Lip)
+        p, pi = _proximal_operator(y, f, m, lambd, Lip)
 
         # COMPUTATION OF THE OBJECTIVE FUNCTION 'Jp' IN 'p'
         tmp = p.flatten()
         Hvp = np.matmul(Hv, tmp)
-        f_Rp = get_entropy(p, m)
+        f_Rp = _get_entropy(p, m)
 
         diff_Vp = Hvp - Visib
         f_0 = (diff_Vp**2).sum()
@@ -491,7 +497,8 @@ def optimise_fb(Hv, Visib, Lip, flux, lambd, shape, pixel, maxiter, tol):
         tau = (t_old - 1.0) / t
         z = x + tau * (x - x_old) + (t_old / t) * (p - x)
 
-        logger.info(f"Iter: {i}, Obj function: {J}")
+        if i % 25 == 0:
+            logger.info(f"Iter: {i}, Obj function: {J}")
 
         if check and (np.sqrt(((x - x_old) ** 2.0).sum()) < tol * np.sqrt((x_old**2).sum())):
             break
@@ -558,7 +565,7 @@ def resistant_mean(data, sigma_cut):
     return mean, sigma
 
 
-def get_percent_lambda(vis):
+def _get_percent_lambda(vis):
     r"""
     Return 'percent_lambda' use with MEM
 
@@ -602,52 +609,63 @@ def get_percent_lambda(vis):
     return percent_lambda
 
 
-def mem(vis, percent_lambda=None, shape=None, pixel=None, maxiter=1000, tol=1e-3, map=True):
+@apu.quantity_input
+def mem(
+    vis: Visibility,
+    shape: Quantity[apu.pix],
+    pixel_size: Quantity[apu.arcsec / apu.pix],
+    *,
+    percent_lambda: Optional[Quantity[apu.percent]] = None,
+    maxiter: int = 1000,
+    tol: float = 1e-3,
+    map: bool = True,
+) -> Union[Quantity, NDArray[np.float64]]:
     r"""
-    Maximum Entropy Method for visibility based image reconstruction
+    Maximum Entropy Method visibility based image reconstruction
 
     Parameters
     ----------
     vis :
         Input Visibilities
-    percent_lambda
-        value used to compute the regularization parameter as a percentage of a maximum value
-        automatically overestimated by the algorithm. Must be in the range [0.0001,0.2]
-    shape
+    shape :
         Image size
-    pixel
+    pixel_size :
         Pixel size
-    maxiter : int
-        Maximum number of iterations of the optimization loop
+    percent_lambda
+        Value used to compute the regularization parameter as a percentage of a maximum value
+        automatically overestimated by the algorithm. Must be in the range [0.0001,0.2]
+    maxiter :
+        Maximum number of iterations of the optimisation loop
     tol : float
         tolerance value used in the stopping rule ( || x - x_old || <= tol || x_old ||)
-
+    map :
+        Return a sunpy map or bare array
     Returns
     -------
 
     """
-    total_flux = estimate_flux(vis, shape, pixel)
+    total_flux = _estimate_flux(vis, shape, pixel_size)
     if percent_lambda is None:
-        percent_lambda = get_percent_lambda(vis)
+        percent_lambda = _get_percent_lambda(vis)
 
-    mean_vis = get_mean_visibilities(vis, shape, pixel)
-    Hv, Lip, Visib = _prepare_for_optimise(pixel, shape, mean_vis)
+    mean_vis = _get_mean_visibilities(vis, shape, pixel_size)
+    Hv, Lip, Visib = _prepare_for_optimise(pixel_size, shape, mean_vis)
 
     # should have same apu as image ct/ keV cm s arcsec**2
     x = np.zeros(shape.to_value("pix").astype(int)).flatten()
 
     # COMPUTATION OF THE; OBJECTIVE; FUNCTION; 'chi2'
-    x = x + total_flux / (shape[0] * shape[1] * pixel[0] * pixel[1]).value
+    x = x + total_flux / (shape[0] * shape[1] * pixel_size[0] * pixel_size[1]).value
     Hvx = np.matmul(Hv, x)
 
     lambd = 2 * np.abs(np.matmul((Hvx.value - Visib), Hv)).max() * percent_lambda
 
-    im = optimise_fb(Hv, Visib, Lip, total_flux, lambd, shape, pixel, maxiter, tol)
+    im = _optimise_fb(Hv, Visib, Lip, total_flux, lambd, shape, pixel_size, maxiter, tol)
 
     # This is needed to match IDL output
     # im = np.rot90(im, -1)
 
     if map:
-        header = generate_header(vis, shape=shape, pixel_size=pixel)
+        header = generate_header(vis, shape=shape, pixel_size=pixel_size)
         return Map((im, header))
     return im
diff --git a/xrayvision/transform.py b/xrayvision/transform.py
index 4ceebcb..d1ed6b7 100644
--- a/xrayvision/transform.py
+++ b/xrayvision/transform.py
@@ -7,128 +7,163 @@
 
 """
 
+from typing import Union, Optional
+
 import astropy.units as apu
 import numpy as np
+import numpy.typing as npt
+from astropy.units import Quantity
 from astropy.units.core import UnitsError
 
+__all__ = ["generate_xy", "generate_uv", "dft_map", "idft_map"]
+
 
-@apu.quantity_input(center="angle", pixel_size="angle")
-def generate_xy(number_pixels, center=0.0 * apu.arcsec, pixel_size=1.0 * apu.arcsec):
+@apu.quantity_input()
+def generate_xy(
+    number_pixels: Quantity[apu.pix],
+    *,
+    phase_centre: Optional[Quantity[apu.arcsec]] = 0.0 * apu.arcsec,
+    pixel_size: Optional[Quantity[apu.arcsec / apu.pix]] = 1.0 * apu.arcsec / apu.pix,
+) -> Quantity[apu.arcsec]:
     """
-    Generate the x or y coordinates given the number of pixels, center and pixel size.
+    Generate the x or y coordinates given the number of pixels, phase_centre and pixel size.
 
     Parameters
     ----------
     number_pixels : `int`
         Number of pixels
-    center : `float`, optional
+    phase_centre : `float`, optional
         Center coordinates
     pixel_size : `float`, optional
         Size of pixel in physical units (e.g. arcsecs, meters)
 
     Returns
     -------
-    `numpy.array`
+    :
         The generated x, y coordinates
 
     See Also
     --------
-    `generate_uv` : Generates corresponding coordinates but un u, v space
+    generate_uv:
+        Generates corresponding coordinates but in Fourier or u, v space.
 
     Examples
     --------
-    >>> generate_xy(9)
+    >>> import astropy.units as apu
+    >>> generate_xy(9*apu.pix)
     <Quantity [-4., -3., -2., -1.,  0.,  1.,  2.,  3.,  4.] arcsec>
 
-    >>> generate_xy(9, pixel_size=2.5 * apu.arcsec)
-    <Quantity [-10. , -7.5, -5. , -2.5,  0. ,  2.5,  5. ,  7.5, 10. ] arcsec>
+    >>> generate_xy(9*apu.pix, pixel_size=2.5 * apu.arcsec/apu.pix)
+    <Quantity [-10. ,  -7.5,  -5. ,  -2.5,   0. ,   2.5,   5. ,   7.5,  10. ] arcsec>
 
-    >>> generate_xy(9, center=10 * apu.arcsec, pixel_size=2.5 * apu.arcsec)
+    >>> generate_xy(9*apu.pix, phase_centre=10 * apu.arcsec, pixel_size=2.5 * apu.arcsec/apu.pix)
     <Quantity [ 0. ,  2.5,  5. ,  7.5, 10. , 12.5, 15. , 17.5, 20. ] arcsec>
 
     """
-    x = (np.arange(number_pixels) - number_pixels / 2 + 0.5) * pixel_size + center
+    x = (
+        np.arange(number_pixels.to_value(apu.pixel)) - (number_pixels.to_value(apu.pix) / 2) + 0.5
+    ) * apu.pix * pixel_size + phase_centre
     return x
 
 
-@apu.quantity_input(center="angle", pixel_size="angle")
-def generate_uv(number_pixels, center=0.0 * apu.arcsec, pixel_size=1.0 * apu.arcsec):
+@apu.quantity_input()
+def generate_uv(
+    number_pixels: Quantity[apu.pix],
+    *,
+    phase_centre: Optional[Quantity[apu.arcsec]] = 0.0 * apu.arcsec,
+    pixel_size: Optional[Quantity[apu.arcsec / apu.pix]] = 1.0 * apu.arcsec / apu.pix,
+) -> Quantity[1 / apu.arcsec]:
     """
-    Generate the u or v  coordinates given the number of pixels, center and pixel size.
+    Generate the u or v coordinates given the number of pixels, phase_centre and pixel size.
 
     Parameters
     ----------
     number_pixels : `int`
         Number of pixels
-    center : `float`, optional
+    phase_centre : `float`, optional
         Center coordinates
     pixel_size : `float`, optional
         Size of pixel in physical units (e.g. arcsecs, meters)
 
     Returns
     -------
-    `numpy.array`
-        The generated u, v coordinates
+    :
+        The generated u, v coordinates.
 
     See Also
     --------
-    `generate_xy` : Generates corresponding coordinate but un x, y space
+    generate_xy:
+        Generates corresponding coordinates but in Fourier or u, v space.
 
     Examples
     --------
-    >>> generate_uv(9)
+    >>> import astropy.units as apu
+    >>> generate_uv(9*apu.pix)
     <Quantity [-0.44444444, -0.33333333, -0.22222222, -0.11111111,  0.        ,
                 0.11111111,  0.22222222,  0.33333333,  0.44444444] 1 / arcsec>
 
-    >>> generate_uv(9, pixel_size=2.5 * apu.arcsec)
+    >>> generate_uv(9*apu.pix, pixel_size=2.5 * apu.arcsec/apu.pix)
     <Quantity [-0.17777778, -0.13333333, -0.08888889, -0.04444444,  0.        ,
                 0.04444444,  0.08888889,  0.13333333,  0.17777778] 1 / arcsec>
 
-    >>> generate_uv(9, center=10 * apu.arcsec, pixel_size=2.5 * apu.arcsec)
+    >>> generate_uv(9*apu.pix, phase_centre=10 * apu.arcsec, pixel_size=2.5 * apu.arcsec/apu.pix)
     <Quantity [-0.07777778, -0.03333333,  0.01111111,  0.05555556,  0.1       ,
                 0.14444444,  0.18888889,  0.23333333,  0.27777778] 1 / arcsec>
 
     """
-    x = (np.arange(number_pixels) - number_pixels / 2 + 0.5) / (pixel_size * number_pixels)
-    if center.value != 0.0:
-        x += 1 / center
+    # x = (np.arange(number_pixels) - number_pixels / 2 + 0.5) / (pixel_size * number_pixels)
+
+    x = (np.arange(number_pixels.to_value(apu.pixel)) - (number_pixels.to_value(apu.pix) / 2) + 0.5) / (
+        pixel_size * number_pixels
+    )
+
+    if phase_centre.value != 0.0:  # type: ignore
+        x += 1 / phase_centre  # type: ignore
     return x
 
 
-@apu.quantity_input(center="angle", pixel_size="angle")
-def dft_map(input_array, *, u, v, center=(0.0, 0.0) * apu.arcsec, pixel_size=(1.0, 1.0) * apu.arcsec):
+@apu.quantity_input()
+def dft_map(
+    input_array: Union[Quantity, npt.NDArray],
+    *,
+    u: Quantity[1 / apu.arcsec],
+    v: Quantity[1 / apu.arcsec],
+    phase_centre: Quantity[apu.arcsec] = (0.0, 0.0) * apu.arcsec,
+    pixel_size: Quantity[apu.arcsec / apu.pix] = (1.0, 1.0) * apu.arcsec / apu.pix,
+) -> Union[Quantity, npt.NDArray]:
     r"""
     Discrete Fourier transform in terms of coordinates returning 1-D array complex visibilities.
 
     Parameters
     ----------
-    input_array : `numpy.ndarray`
+    input_array :
         Input array to be transformed should be 2D (m, n)
-    uv : `numpy.array`
-        Array of 2xN u, v coordinates where the visibilities will be evaluated
-    center : `float` (x, y), optional
-        Coordinates of the center of the map e.g. ``(0,0)`` or ``[5.0, -2.0]``
+    u :
+        Array of 2xN u coordinates where the visibilities are evaluated.
+    v :
+        Array of 2xN v coordinates where the visibilities are evaluated.
+    phase_centre :
+        Coordinates of the phase_centre of the map e.g. ``(0,0)`` or ``[5.0, -2.0]``.
     pixel_size : `float` (dx, dy), optional
-        The pixel size in x and y directions, need not be square e.g. ``(1, 3)``
+        The pixel size need not be square e.g. ``(1, 3)``.
 
     Returns
     -------
-    `numpy.ndarray`
-        Array of N `complex` visibilities evaluated at the u, v coordinates given by `uv`
+    :
+        Array of N `complex` visibilities evaluated at the given `u`, `v` coordinates.
 
     """
-    m, n = input_array.shape
+    m, n = input_array.shape * apu.pix
+    x = generate_xy(m, phase_centre=phase_centre[0], pixel_size=pixel_size[0])  # type: ignore
+    y = generate_xy(n, phase_centre=phase_centre[1], pixel_size=pixel_size[1])  # type: ignore
 
-    y = generate_xy(m, center[1], pixel_size[1])
-    x = generate_xy(n, center[0], pixel_size[0])
-
-    x, y = np.meshgrid(x, y)
+    x, y = np.meshgrid(x, y, indexing="ij")
     uv = np.vstack([u, v])
     # Check units are correct for exp need to be dimensionless and then remove units for speed
     if (uv[0, :] * x[0, 0]).unit == apu.dimensionless_unscaled and (
         uv[1, :] * y[0, 0]
     ).unit == apu.dimensionless_unscaled:
-        uv = uv.value
+        uv = uv.value  # type: ignore
         x = x.value
         y = y.value
 
@@ -147,38 +182,48 @@ def dft_map(input_array, *, u, v, center=(0.0, 0.0) * apu.arcsec, pixel_size=(1.
         )
 
 
-@apu.quantity_input(center="angle", pixel_size="angle")
+@apu.quantity_input
 def idft_map(
-    input_vis, *, u, v, shape, weights=None, center=(0.0, 0.0) * apu.arcsec, pixel_size=(1.0, 1.0) * apu.arcsec
-):
+    input_vis: Union[Quantity, npt.NDArray],
+    *,
+    u: Quantity[1 / apu.arcsec],
+    v: Quantity[1 / apu.arcsec],
+    shape: Quantity[apu.pix],
+    weights: Optional[npt.NDArray] = None,
+    phase_centre: Quantity[apu.arcsec] = (0.0, 0.0) * apu.arcsec,
+    pixel_size: Quantity[apu.arcsec / apu.pix] = (1.0, 1.0) * apu.arcsec / apu.pix,
+) -> Union[Quantity, npt.NDArray]:
     r"""
     Inverse discrete Fourier transform in terms of coordinates returning a 2D real array or image.
 
     Parameters
     ----------
-    uv : `numpy.ndarray`
-        Array of 2xN u, v coordinates corresponding to the input visibilities in `input_vis`
-    input_vis : `numpy.ndarray`
-        Array of N `complex` input visibilities
-    shape : `float` (m,n)
-        The shape of the output array to create
-    weights : `numpy.ndarray`
-        Array of weights for visibilities
-    center : `float` (x, y), optional
-        Coordinates of the center of the map e.g. ``(0,0)`` or ``[5.0, -2.0]``
-    pixel_size : `float` (dx, dy), optional
-        The pixel size in x and y directions, need not be square e.g. ``(1, 3)``
+    input_vis :
+        Input array of N complex visibilities to be transformed to a 2D array.
+    u :
+        Array of N u coordinates corresponding to the input visibilities in `input_vis`
+    v :
+        Array of N v coordinates corresponding to the input visibilities in `input_vis`
+    shape :
+        The shape of the output array to create.
+    weights :
+        Array of weights for visibilities.
+    phase_centre :
+        Coordinates of the phase_centre of the map e.g. ``(0,0)`` or ``[5.0, -2.0]``.
+    pixel_size :
+        The pixel size this need not be square e.g. ``(1, 3)``.
 
     Returns
     -------
-    `numpy.ndarray`
-        The complex visibilities evaluated at the u, v coordinates
+    :
+        2D image obtained from the visibilities evaluated at the given `u`, `v` coordinates.
 
     """
     m, n = shape
-    y = generate_xy(m, center[1], pixel_size[1])
-    x = generate_xy(n, center[0], pixel_size[0])
-    x, y = np.meshgrid(x, y)
+    x = generate_xy(m, phase_centre=phase_centre[0], pixel_size=pixel_size[0])  # type: ignore
+    y = generate_xy(n, phase_centre=phase_centre[1], pixel_size=pixel_size[1])  # type: ignore
+
+    x, y = np.meshgrid(x, y, indexing="ij")
 
     if weights is None:
         weights = np.ones(input_vis.shape)
@@ -187,7 +232,7 @@ def idft_map(
     if (uv[0, :] * x[0, 0]).unit == apu.dimensionless_unscaled and (
         uv[1, :] * y[0, 0]
     ).unit == apu.dimensionless_unscaled:
-        uv = uv.value
+        uv = uv.value  # type: ignore
         x = x.value
         y = y.value
 
diff --git a/xrayvision/utils.py b/xrayvision/utils.py
index 2935b24..bb9789f 100644
--- a/xrayvision/utils.py
+++ b/xrayvision/utils.py
@@ -1,7 +1,8 @@
 import logging
+from typing import Union, Optional
 
 
-def get_logger(name, level=logging.WARNING):
+def get_logger(name: str, level: Optional[Union[int, str]] = logging.WARNING) -> logging.Logger:
     """
     Return a configured logger instance.
 
@@ -18,9 +19,9 @@ def get_logger(name, level=logging.WARNING):
         Configured logger
     """
     logger = logging.getLogger(name)
-    logger.setLevel(level)
+    logger.setLevel(level)  # type: ignore
     handler = logging.StreamHandler()
-    handler.setLevel(level)
+    handler.setLevel(level)  # type: ignore
     formatter = logging.Formatter(
         "%(asctime)s %(levelname)s %(name)s %(lineno)s: %(message)s", datefmt="%Y-%m-%dT%H:%M:%SZ"
     )
diff --git a/xrayvision/visibility.py b/xrayvision/visibility.py
index 2d4fb75..66a5398 100644
--- a/xrayvision/visibility.py
+++ b/xrayvision/visibility.py
@@ -1,5 +1,5 @@
 """
-Modules contains visibility related classes.
+Module contains visibility related classes.
 
 This contains classes to hold general visibilities and specialised classes hold visibilities from
 certain spacecraft or instruments
@@ -7,7 +7,7 @@
 
 import abc
 from typing import Union, Optional
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 
 import astropy.units as apu
 import numpy as np
@@ -17,6 +17,8 @@
 
 __all__ = ["Visibility", "Visibilities", "VisMeta", "VisibilitiesABC", "VisMetaABC"]
 
+from astropy.units import Quantity
+
 _E_RANGE_KEY = "spectral_range"
 _T_RANGE_KEY = "time_range"
 _OBS_COORD_KEY = "observer_coordinate"
@@ -34,21 +36,21 @@ def observer_coordinate(self) -> Union[SkyCoord, None]:
 
     @property
     @abc.abstractmethod
-    def spectral_range(self) -> Union[Iterable[apu.Quantity], None]:
+    def spectral_range(self) -> Optional[Iterable[apu.Quantity]]:
         """
         Spectral range over which the visibilities are computed.
         """
 
     @property
     @abc.abstractmethod
-    def time_range(self) -> Union[Iterable[Time], None]:
+    def time_range(self) -> Optional[Iterable[Time]]:
         """
         Time range over which the visibilities are computed.
         """
 
     @property
     @abc.abstractmethod
-    def vis_labels(self) -> Union[Iterable[str], None]:
+    def vis_labels(self) -> Sequence[Iterable[str]]:
         """
         Labels of each visibility.
         """
@@ -133,6 +135,60 @@ def phase_uncertainty(self) -> Union[Iterable[apu.Quantity[apu.deg]], None]:
         """
 
 
+class VisMeta(VisMetaABC, dict):
+    """
+    A class for holding Visibility-specific metadata.
+
+    Parameters
+    ----------
+    meta: `dict`
+        A dictionary of the metadata
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Check controlled/expected inputs are of correct type and units.
+        controled_args = (
+            (_OBS_COORD_KEY, SkyCoord),
+            (_E_RANGE_KEY, apu.Quantity, apu.keV, apu.spectral()),
+            (_T_RANGE_KEY, Time),
+        )
+        for args in controled_args:
+            self._check_input_type_and_unit(*args)
+
+    def _check_input_type_and_unit(self, key, key_type, unit=None, equivalencies=None):
+        value = self.get(key, None)
+        if not isinstance(value, (key_type, type(None))):
+            raise KeyError(f"Inputs must include a key, '{key}', that gives a {key_type}.")
+        if unit is not None and value is not None and not value.unit.is_equivalent(unit, equivalencies=equivalencies):
+            raise ValueError(f"'{key}' must have angular units.")
+
+    @property
+    def observer_coordinate(self):
+        return self.get(_OBS_COORD_KEY, None)
+
+    @property
+    def spectral_range(self):
+        return self.get(_E_RANGE_KEY, None)
+
+    @property
+    def time_range(self):
+        return self.get(_T_RANGE_KEY, None)
+
+    @property
+    def vis_labels(self):
+        return self.get(_VIS_LABELS_KEY, None)
+
+    @property
+    def instrument(self):
+        instr = None
+        i, n = 0, len(_INSTR_KEYS)
+        while not instr and i < n:
+            instr = self.get(_INSTR_KEYS[i], None)
+            i += 1
+        return instr
+
+
 class Visibilities(VisibilitiesABC):
     @apu.quantity_input()
     def __init__(
@@ -224,6 +280,10 @@ def __init__(
         self._uv_key = "uv"
         self._units_key = "units"
 
+        # Build meta. Make sure that phase center is included.
+        if not isinstance(meta, VisMetaABC):
+            meta = VisMeta(meta)
+
         # Construct underlying data object.
         # In case visibilities is multi-dimensional, assume last axis is the uv-axis.
         # and give other axes arbitrary names.
@@ -244,12 +304,12 @@ def __init__(
             data[self._phase_uncert_key] = (dims, phase_uncertainty.to_value(phase.unit))
         if meta is None:
             meta = VisMeta(dict())
-        vis_labels = meta.vis_labels
+        vis_labels = getattr(meta, "vis_labels", None)
         if vis_labels is not None:
             if len(vis_labels) != nvis:
                 raise ValueError(
                     "meta.vis_labels must be same length as number of visibilites. "
-                    f"Number of labels = {len(meta.vis_labels)}; "
+                    f"Number of labels = {len(vis_labels)}; "
                     f"Number of visibilities = {nvis}"
                 )
             coords[_VIS_LABELS_KEY] = ([self._uv_key], vis_labels)
@@ -353,113 +413,48 @@ def __repr__(self):
         return f"{self.__class__.__name__}< {self.u.size}, {self.visibilities}>"
 
 
-class VisMeta(VisMetaABC, dict):
-    """
-    A class for holding Visibility-specific metadata.
-
-    Parameters
-    ----------
-    meta: `dict`
-        A dictionary of the metadata
-    """
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        # Check controlled/expected inputs are of correct type and units.
-        controled_args = (
-            (_OBS_COORD_KEY, SkyCoord),
-            (_E_RANGE_KEY, apu.Quantity, apu.keV, apu.spectral()),
-            (_T_RANGE_KEY, Time),
-        )
-        for args in controled_args:
-            self._check_input_type_and_unit(*args)
-
-    def _check_input_type_and_unit(self, key, key_type, unit=None, equivalencies=None):
-        value = self.get(key, None)
-        if not isinstance(value, (key_type, type(None))):
-            raise KeyError(f"Inputs must include a key, '{key}', that gives a {key_type}.")
-        if unit is not None and value is not None and not value.unit.is_equivalent(unit, equivalencies=equivalencies):
-            raise ValueError(f"'{key}' must have angular units.")
-
-    @property
-    def observer_coordinate(self):
-        return self.get(_OBS_COORD_KEY, None)
-
-    @property
-    def spectral_range(self):
-        return self.get(_E_RANGE_KEY, None)
-
-    @property
-    def time_range(self):
-        return self.get(_T_RANGE_KEY, None)
-
-    @property
-    def vis_labels(self):
-        return self.get(_VIS_LABELS_KEY, None)
-
-    @property
-    def instrument(self):
-        instr = None
-        i, n = 0, len(_INSTR_KEYS)
-        while not instr and i < n:
-            instr = self.get(_INSTR_KEYS[i], None)
-            i += 1
-        return instr
-
-
-class BaseVisibility:
-    r"""
-    Base visibility containing bare essential fields, u, v, and complex vis
-    """
-
-    @apu.quantity_input(u=1 / apu.arcsec, v=1 / apu.arcsec, center=apu.arcsec)
-    def __int__(self, u, v, vis, center=(0, 0) * apu.arcsec):
-        self.u = u
-        self.v = v
-        self.vis = vis
-        self.center = center
-
-
 class Visibility:
     r"""
     Hold a set of related visibilities and information.
 
-    Attributes
-    ----------
-    vis : `numpy.ndarray`
-        Array of N complex visibilities at coordinates in `uv`
-    u : `numpy.ndarray`
-        Array of `u` coordinates where visibilities will be evaluated
-    v : `numpy.ndarray`
-        Array of `v` coordinates where visibilities will be evaluated
-    center : `float` (x, y), optional
-        The x, y offset of phase center
 
     """
 
-    @apu.quantity_input(uv=1 / apu.arcsec, offset=apu.arcsec, center=apu.arcsec, pixel_size=apu.arcsec)
-    def __init__(self, vis, *, u, v, offset=(0.0, 0.0) * apu.arcsec, center=(0.0, 0.0) * apu.arcsec):
+    @apu.quantity_input
+    def __init__(
+        self,
+        vis,
+        *,
+        u: Quantity[1 / apu.arcsec],
+        v: Quantity[1 / apu.arcsec],
+        offset: Optional[Quantity[apu.arcsec]] = (0.0, 0.0) * apu.arcsec,
+        phase_centre: Optional[Quantity[apu.arcsec]] = (0.0, 0.0) * apu.arcsec,
+    ) -> None:
         r"""
-        Initialise a new Visibility object.
+        Generic Visibility object.
 
         Parameters
         ----------
-        vis : `numpy.ndarray`
-            Array of N complex visibilities at coordinates in `uv`.
-        u : `numpy.ndarray`
-            Array of `u` coordinates where visibilities will be evaluated.
-        v : `numpy.ndarray`
-            Array of `v` coordinates where visibilities will be evaluated.
-        center :
-            Phase centre
-        """
-        self.u = u
-        self.v = v
-        self.vis = vis
-        self.center = center
-        self.offset = offset
-
-    def __repr__(self):
+        vis:
+            Array of N complex visibilities sampled at the `u`, `v` coordinates.
+        u:
+            Array of `u` coordinates where visibilities are sampled.
+        v:
+            Array of `v` coordinates where visibilities are sampled.
+        phase_centre:
+            Phase centre of the visibility, defaults to (0,0).
+        offset:
+            Offset of the phase_centre visibility, defaults to (0,0).
+
+
+        """
+        self.u: Quantity[1 / apu.arcsec] = u
+        self.v: Quantity[1 / apu.arcsec] = v
+        self.vis: Quantity = vis
+        self.phase_centre: Quantity[apu.arcsec] = phase_centre
+        self.offset: Quantity[apu.arcsec] = offset
+
+    def __repr__(self) -> str:
         r"""
         Return a printable representation of the visibility.
 
@@ -470,7 +465,7 @@ def __repr__(self):
         """
         return f"{self.__class__.__name__}< {self.u.size}, {self.vis}>"
 
-    def __eq__(self, other):
+    def __eq__(self, other) -> bool:
         r"""
         Equality for Visibility class
 
@@ -490,5 +485,5 @@ def __eq__(self, other):
 
         if all(props_equal):
             return True
-        else:
-            return False
+
+        return False


Test patch:
diff --git a/xrayvision/tests/test_clean.py b/xrayvision/tests/test_clean.py
index 31e781c..0e9673c 100644
--- a/xrayvision/tests/test_clean.py
+++ b/xrayvision/tests/test_clean.py
@@ -1,14 +1,15 @@
 import astropy.units as u
 import numpy as np
 from astropy.convolution.kernels import Gaussian2DKernel
+from numpy.testing import assert_allclose
 from scipy import signal
 
 from xrayvision.clean import (
+    _component,
+    _radial_prolate_sphereoidal,
+    _vec_radial_prolate_sphereoidal,
     clean,
-    component,
     ms_clean,
-    radial_prolate_sphereoidal,
-    vec_radial_prolate_sphereoidal,
     vis_clean,
 )
 from xrayvision.imaging import image_to_vis
@@ -23,6 +24,7 @@ def test_clean_ideal():
     clean_map = np.zeros((n, m))
     clean_map[pos1[0], pos1[1]] = 10.0
     clean_map[pos2[0], pos2[1]] = 7.0
+    clean_map = clean_map  # << 1 / u.arcsec**2
 
     dirty_beam = np.zeros((n, m))
     dirty_beam[(n - 1) // 4 : (n - 1) // 4 + (n - 1) // 2, (m - 1) // 2] = 0.75
@@ -36,35 +38,35 @@ def test_clean_ideal():
     dirty_map = signal.convolve(clean_map, dirty_beam, mode="same")
 
     # Disable convolution of model with gaussian for testing
-    out_map, model, resid = clean(dirty_map, dirty_beam, clean_beam_width=0.0)
+    out_map, model, resid = clean(dirty_map, dirty_beam, clean_beam_width=None)
 
     # Within threshold default threshold of 0.1
-    assert np.allclose(clean_map, out_map, atol=dirty_beam.max() * 0.1)
+    assert_allclose(out_map, clean_map, atol=dirty_beam.max() * 1e-13)
 
 
 def test_component():
     comp = np.zeros((3, 3))
     comp[1, 1] = 1.0
 
-    res = component(scale=0, shape=(3, 3))
+    res = _component(scale=0, shape=(3, 3))
     assert np.array_equal(res, comp)
 
-    res = component(scale=1, shape=(3, 3))
+    res = _component(scale=1, shape=(3, 3))
     assert np.array_equal(res, comp)
 
-    res = component(scale=2, shape=(6, 6))
+    res = _component(scale=2, shape=(6, 6))
     assert np.all(res[0, :] == 0.0)
     assert np.all(res[:, 0] == 0.0)
     assert np.all(res[2:4, 2:4] == res.max())
 
-    res = component(scale=3, shape=(7, 7))
+    res = _component(scale=3, shape=(7, 7))
     assert np.all(res[0, :] == 0.0)
     assert np.all(res[:, 0] == 0.0)
     assert res[3, 3] == 1
 
 
 def test_radial_prolate_spheroidal():
-    amps = [radial_prolate_sphereoidal(r) for r in [-1.0, 0.0, 0.5, 1.0, 2.0]]
+    amps = [_radial_prolate_sphereoidal(r) for r in [-1.0, 0.0, 0.5, 1.0, 2.0]]
     assert amps[0] == 1.0
     assert amps[1] == 1.0
     assert amps[2] == 0.36106538453111797
@@ -74,8 +76,8 @@ def test_radial_prolate_spheroidal():
 
 def test_vec_radial_prolate_spheroidal():
     radii = np.linspace(-0.5, 1.5, 1000)
-    amps1 = [radial_prolate_sphereoidal(r) for r in radii]
-    amps2 = vec_radial_prolate_sphereoidal(radii)
+    amps1 = [_radial_prolate_sphereoidal(r) for r in radii]
+    amps2 = _vec_radial_prolate_sphereoidal(radii)
     assert np.allclose(amps1, amps2)
 
 
@@ -100,7 +102,9 @@ def test_ms_clean_ideal():
     dirty_map = signal.convolve2d(clean_map, dirty_beam, mode="same")
 
     # Disable convolution of model with gaussian for testing
-    model, res = ms_clean(dirty_map, dirty_beam, [1, 1] * u.arcsec, scales=[1], clean_beam_width=0.0)
+    model, res = ms_clean(
+        dirty_map, dirty_beam, pixel_size=[1, 1] * u.arcsec / u.pixel, scales=[1], clean_beam_width=None
+    )
     recovered = model + res
 
     # Within threshold default threshold
@@ -129,16 +133,16 @@ def test_clean_sim():
     sub_uv = np.hstack([sub_uv, np.zeros((2, 1))]) / u.arcsec
 
     # Factor of 9 is compensate for the factor of  3 * 3 increase in size
-    dirty_beam = idft_map(np.ones(321) / 321, u=sub_uv[0, :], v=sub_uv[1, :], shape=(n * 3 + 1, m * 3 + 1))
+    dirty_beam = idft_map(np.ones(321) / 321, u=sub_uv[0, :], v=sub_uv[1, :], shape=(n * 3 + 1, m * 3 + 1) * u.pix)
 
     vis = dft_map(data, u=sub_uv[0, :], v=sub_uv[1, :])
 
-    dirty_map = idft_map(vis, u=sub_uv[0, :], v=sub_uv[1, :], shape=(n, m))
+    dirty_map = idft_map(vis, weights=1 / 321, u=sub_uv[0, :], v=sub_uv[1, :], shape=(n, m) * u.pix)
 
     clean_map, model, res = clean(
-        dirty_map, dirty_beam, clean_beam_width=0.1 * u.arcsec, pixel=[2, 2] * u.arcsec, niter=500
+        dirty_map, dirty_beam, pixel_size=[2, 2] * u.arcsec / u.pix, clean_beam_width=0.1 * u.arcsec, niter=500
     )
-    np.allclose(data, clean_map.data, rtol=dirty_beam.max() * 0.1)
+    assert_allclose(clean_map, data, atol=dirty_beam.max() * 0.1)
 
 
 def test_vis_clean_sim():
@@ -164,6 +168,11 @@ def test_vis_clean_sim():
     vis = image_to_vis(data * u.dimensionless_unscaled, u=sub_uv[0, :], v=sub_uv[1, :])
 
     clean_map, model, res = vis_clean(
-        vis, shape=(m, n) * u.pix, clean_beam_width=0, pixel=[2, 2] * u.arcsec, natural=False, niter=100
+        vis,
+        shape=(m, n) * u.pix,
+        pixel_size=[2, 2] * u.arcsec / u.pix,
+        clean_beam_width=None,
+        niter=100,
+        scheme="uniform",
     )
     np.allclose(data, clean_map.data, atol=0.1)
diff --git a/xrayvision/tests/test_imaging.py b/xrayvision/tests/test_imaging.py
index 6ea4b98..82819c4 100644
--- a/xrayvision/tests/test_imaging.py
+++ b/xrayvision/tests/test_imaging.py
@@ -2,6 +2,7 @@
 import numpy as np
 import pytest
 from astropy.convolution.kernels import Gaussian2DKernel
+from numpy.testing import assert_allclose
 from sunpy.map import Map
 
 from xrayvision.imaging import image_to_vis, map_to_vis, vis_psf_image, vis_to_image, vis_to_map
@@ -27,30 +28,88 @@ def uv():
     return u.flatten() / apu.arcsec, v.flatten() / apu.arcsec
 
 
+@pytest.mark.parametrize("shape", [(10, 10), (5, 5), (2, 2)])
+def test_image_to_vis_vis_to_image(shape):
+    # Images with the same total flux
+    extent = [10, 10] * apu.arcsec
+    arr1 = np.full((shape), 1) * apu.ct / np.prod(shape)
+    ps1 = extent / (shape * apu.pix)
+    img1 = arr1 / (ps1[0] * ps1[1] * apu.pix**2)
+
+    # flux in image is 1
+    assert_allclose(img1.sum() * ps1[0] * apu.pix * ps1[1] * apu.pix, 1 * apu.ct)
+
+    # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT)
+    u = generate_uv((shape[0]) * apu.pix, pixel_size=ps1[0])
+    v = generate_uv((shape[1]) * apu.pix, pixel_size=ps1[1])
+    u, v = np.meshgrid(u, v, indexing="ij")
+    u, v = np.array([u, v]).reshape(2, np.prod(shape)) / apu.arcsec
+
+    vis1 = image_to_vis(img1, u=u, v=v, pixel_size=ps1)
+    res1 = vis_to_image(vis1, shape=shape * apu.pix, pixel_size=ps1)
+    assert_allclose(res1, img1)
+
+
+@pytest.mark.skip(reason="WIP")
+def test_vis_to_image_conserve():
+    shape = (3, 3)
+    shape1 = (5, 5)
+    shape2 = (7, 7)
+    # Images with the same total flux
+    extent = [10, 10] * apu.arcsec
+    arr1 = np.zeros(shape)
+    arr1[shape[0] // 2, shape[1] // 2] = 1 / np.prod(shape)
+    ps = extent / (shape * apu.pix)
+    img1 = arr1 / (ps[0] * ps[1] * apu.pix**2)
+
+    # flux in image is same
+    assert_allclose(img1.sum() * ps[0] * apu.pix * ps[1] * apu.pix, arr1.sum())
+
+    u = generate_uv(shape[0] * apu.pix, pixel_size=ps[0])
+    v = generate_uv(shape[1] * apu.pix, pixel_size=ps[1])
+    u, v = np.meshgrid(u, v, indexing="ij")
+    u, v = np.array([u, v]).reshape(2, np.prod(shape)) / apu.arcsec
+
+    vis1 = image_to_vis(img1, u=u, v=v, pixel_size=ps)
+
+    ps2 = extent / (shape1 * apu.pix)
+    ps3 = extent / (shape2 * apu.pix)
+    res = vis_to_image(vis1, shape=shape * apu.pix, pixel_size=ps)
+    res2 = vis_to_image(vis1, shape=shape1 * apu.pix, pixel_size=ps2)
+    res3 = vis_to_image(vis1, shape=shape2 * apu.pix, pixel_size=ps3)
+
+    assert_allclose(res, img1)
+    assert_allclose(res2, img1)
+    assert_allclose(res3, img1)
+
+
 @pytest.mark.parametrize("pixel_size", [(0.5), (1), (2)])
-def test_psf_to_image(pixel_size, uv):
+def test_vis_to_psf(pixel_size, uv):
     u, v = uv
-    img = np.zeros((65, 65)) * (apu.ph / apu.arcsec**2)
-    img[32, 32] = 1.0 * (apu.ph / apu.arcsec**2)
-    obs_vis = dft_map(img, u=u, v=v, pixel_size=[2.0, 2.0] * apu.arcsec)
+    ps = [pixel_size, pixel_size] * apu.arcsec / apu.pix
+    img = np.zeros((65, 65)) * (apu.ph / apu.cm**2)
+    img[32, 32] = 1 * (apu.ph / apu.cm**2)  # pixel size of 4 -> 1/4
+    obs_vis = dft_map(img, u=u, v=v, pixel_size=ps)
     weights = np.sqrt(u**2 + v**2).value
     weights /= weights.sum()
-    psf_calc = idft_map(obs_vis, u=u, v=v, shape=[65, 65], pixel_size=[2, 2] * apu.arcsec, weights=weights)
+    psf_calc = idft_map(obs_vis, u=u, v=v, shape=[65, 65] * apu.pix, pixel_size=ps, weights=weights)
     vis = Visibility(obs_vis, u=u, v=v)
-    res = vis_psf_image(vis, shape=[65, 65] * apu.pixel, pixel_size=2 * apu.arcsec, natural=False)
-    assert np.allclose(psf_calc, res)
+    res = vis_psf_image(vis, shape=[65, 65] * apu.pixel, pixel_size=ps, scheme="uniform")
+    assert_allclose(res, psf_calc)
 
 
-def test_vis_to_image(uv):
+def test_vis_to_image_against_idft(uv):
     u, v = uv
-    img = np.zeros((65, 65)) * (apu.ph / apu.arcsec**2)
-    img[32, 32] = 1.0 * (apu.ph / apu.arcsec**2)
-    obs_vis = dft_map(img, u=u, v=v, pixel_size=[2.0, 2.0] * apu.arcsec)
+    img = np.zeros((65, 65)) * (apu.ph / apu.cm**2)
+    img[32, 32] = 1.0 * (apu.ph / apu.cm**2)
+    obs_vis = dft_map(img, u=u, v=v, pixel_size=[2.0, 2.0] * apu.arcsec / apu.pix)
     weights = np.sqrt(u**2 + v**2).value
     weights /= weights.sum()
-    bp_calc = idft_map(obs_vis, u=u, v=v, shape=[65, 65], pixel_size=[2, 2] * apu.arcsec, weights=weights)
+    bp_calc = idft_map(
+        obs_vis, u=u, v=v, shape=[65, 65] * apu.pix, pixel_size=[2, 2] * apu.arcsec / apu.pix, weights=weights
+    )
     vis = Visibility(obs_vis, u=u, v=v)
-    res = vis_to_image(vis, shape=[65, 65] * apu.pixel, pixel_size=2 * apu.arcsec, natural=False)
+    res = vis_to_image(vis, shape=[65, 65] * apu.pixel, pixel_size=2 * apu.arcsec / apu.pix, scheme="uniform")
     assert np.allclose(bp_calc, res)
 
 
@@ -61,14 +120,14 @@ def test_image_to_vis():
     image = np.zeros((m, n))
 
     # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT)
-    u = generate_uv(m)
-    v = generate_uv(n)
-    u, v = np.meshgrid(u, v)
+    u = generate_uv(m * apu.pix)
+    v = generate_uv(n * apu.pix)
+    u, v = np.meshgrid(u, v, indexing="ij")
     u, v = np.array([u, v]).reshape(2, size) / apu.arcsec
 
     # For an empty map visibilities should all be zero (0+0j)
     empty_vis = image_to_vis(image, u=v, v=v)
-    assert np.array_equal(empty_vis.center, (0.0, 0.0) * apu.arcsec)
+    assert np.array_equal(empty_vis.phase_centre, (0.0, 0.0) * apu.arcsec)
     assert np.array_equal(empty_vis.vis, np.zeros(n * m, dtype=complex))
 
 
@@ -79,14 +138,14 @@ def test_image_to_vis_center():
     image = np.zeros((m, n))
 
     # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT)
-    u = generate_uv(m)
-    v = generate_uv(n)
-    u, v = np.meshgrid(u, v)
+    u = generate_uv(m * apu.pix)
+    v = generate_uv(n * apu.pix)
+    u, v = np.meshgrid(u, v, indexing="ij")
     u, v = np.array([u, v]).reshape(2, size) / apu.arcsec
 
     # For an empty map visibilities should all be zero (0+0j)
-    empty_vis = image_to_vis(image, u=u, v=v, center=(2.0, -3.0) * apu.arcsec)
-    assert np.array_equal(empty_vis.center, (2.0, -3.0) * apu.arcsec)
+    empty_vis = image_to_vis(image, u=u, v=v, phase_centre=(2.0, -3.0) * apu.arcsec)
+    assert np.array_equal(empty_vis.offset, (2.0, -3.0) * apu.arcsec)
     assert np.array_equal(empty_vis.vis, np.zeros(n * m, dtype=complex))
 
 
@@ -99,19 +158,19 @@ def test_map_to_vis(pos, pixel):
     size = m * n
 
     pos = pos * apu.arcsec
-    pixel = pixel * apu.arcsec
+    pixel = pixel * apu.arcsec / apu.pix
 
     # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT)
-    u = generate_uv(m, pos[0])
-    v = generate_uv(n, pos[1])
-    u, v = np.meshgrid(u, v)
+    u = generate_uv(m * apu.pix, phase_centre=pos[0], pixel_size=pixel[0])
+    v = generate_uv(n * apu.pix, phase_centre=pos[1], pixel_size=pixel[1])
+    u, v = np.meshgrid(u, v, indexing="ij")
     u, v = np.array([u, v]).reshape(2, size) / apu.arcsec
 
     header = {
-        "crval1": pos[0].value,
-        "crval2": pos[1].value,
-        "cdelt1": pixel[0].value,
-        "cdelt2": pixel[1].value,
+        "crval1": pos[1].value,
+        "crval2": pos[0].value,
+        "cdelt1": pixel[1].value,
+        "cdelt2": pixel[0].value,
         "cunit1": "arcsec",
         "cunit2": "arcsec",
     }
@@ -127,19 +186,21 @@ def test_map_to_vis(pos, pixel):
     assert np.allclose(res, data)
 
 
-def test_vis_to_image1():
-    m = n = 33
+def test_vis_to_image():
+    m = n = 30
     size = m * n
 
     # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT)
-    u = generate_uv(m)
-    v = generate_uv(n)
+    u = generate_uv(m * apu.pix)
+    v = generate_uv(n * apu.pix)
     u, v = np.meshgrid(u, v)
     uv = np.array([u, v]).reshape(2, size) / apu.arcsec
     u, v = uv
 
     # Astropy index order is opposite to that of numpy, is 1st dim is across second down
     data = Gaussian2DKernel(6, x_size=n, y_size=m).array
+    data = np.zeros((m, n))
+    data[m // 2, n // 2] = 1
 
     vis = image_to_vis(data, u=u, v=v)
     res = vis_to_image(vis, shape=(m, n) * apu.pixel)
@@ -152,16 +213,16 @@ def test_vis_to_image_single_pixel_size():
     size = m * n
 
     # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT)
-    u = generate_uv(m, pixel_size=2.0 * apu.arcsec)
-    v = generate_uv(n, pixel_size=2.0 * apu.arcsec)
-    u, v = np.meshgrid(u, v)
+    u = generate_uv(m * apu.pix, pixel_size=2.0 * apu.arcsec / apu.pix)
+    v = generate_uv(n * apu.pix, pixel_size=2.0 * apu.arcsec / apu.pix)
+    u, v = np.meshgrid(u, v, indexing="ij")
     uv = np.array([u, v]).reshape(2, size) / apu.arcsec
     u, v = uv
     # Astropy index order is opposite to that of numpy, is 1st dim is across second down
     data = Gaussian2DKernel(6, x_size=n, y_size=m).array
 
-    vis = image_to_vis(data, u=u, v=v, pixel_size=(2.0, 2.0) * apu.arcsec)
-    res = vis_to_image(vis, shape=(m, n) * apu.pixel, pixel_size=2.0 * apu.arcsec)
+    vis = image_to_vis(data, u=u, v=v, pixel_size=(2.0, 2.0) * apu.arcsec / apu.pix)
+    res = vis_to_image(vis, shape=(m, n) * apu.pixel, pixel_size=2.0 * apu.arcsec / apu.pix)
     assert res.shape == (m, n)
     assert np.allclose(data, res)
 
@@ -171,9 +232,9 @@ def test_vis_to_image_invalid_pixel_size():
     size = m * n
 
     # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT)
-    u = generate_uv(m)
-    v = generate_uv(n)
-    u, v = np.meshgrid(u, v)
+    u = generate_uv(m * apu.pix)
+    v = generate_uv(n * apu.pix)
+    u, v = np.meshgrid(u, v, indexing="ij")
     uv = np.array([u, v]).reshape(2, size) / apu.arcsec
     u, v = uv
 
@@ -182,24 +243,30 @@ def test_vis_to_image_invalid_pixel_size():
 
     vis = image_to_vis(data, u=u, v=v)
     with pytest.raises(ValueError):
-        vis_to_image(vis, shape=(m, n) * apu.pixel, pixel_size=[1, 2, 2] * apu.arcsec)
+        vis_to_image(vis, shape=(m, n) * apu.pixel, pixel_size=[1, 2, 2] * apu.arcsec / apu.pix)
 
 
-@pytest.mark.parametrize("m,n,pos,pixel", [(33, 33, (10.0, -5.0), (2.0, 3.0)), (32, 32, (-12, -19), (1.0, 5.0))])
+@pytest.mark.parametrize(
+    "m,n,pos,pixel",
+    [  # (33, 33, (0, 0), (1.5, 1.5)),
+        (33, 33, (10.0, -5.0), (2.0, 3.0)),
+        (32, 32, (-12, -19), (1.0, 5.0)),
+    ],
+)
 def test_vis_to_map(m, n, pos, pixel):
     pos = pos * apu.arcsec
-    pixel = pixel * apu.arcsec
-    u = generate_uv(m, pixel[0])
-    v = generate_uv(n, pixel[1])
-    u, v = np.meshgrid(u, v)
+    pixel = pixel * apu.arcsec / apu.pix
+    u = generate_uv(m * apu.pix, phase_centre=pos[0], pixel_size=pixel[0])
+    v = generate_uv(n * apu.pix, phase_centre=pos[1], pixel_size=pixel[1])
+    u, v = np.meshgrid(u, v, indexing="ij")
     uv = np.array([u, v]).reshape(2, m * n) / apu.arcsec
     u, v = uv
 
     header = {
-        "crval1": pos[0].value,
-        "crval2": pos[1].value,
-        "cdelt1": pixel[0].value,
-        "cdelt2": pixel[1].value,
+        "crval1": pos[1].value,
+        "crval2": pos[0].value,
+        "cdelt1": pixel[1].value,
+        "cdelt2": pixel[0].value,
         "cunit1": "arcsec",
         "cunit2": "arcsec",
     }
@@ -208,22 +275,23 @@ def test_vis_to_map(m, n, pos, pixel):
 
     vis = map_to_vis(mp, u=u, v=v)
 
-    res = vis_to_map(vis, shape=(m, n) * apu.pixel, pixel_size=pixel, natural=False)
-    # assert np.allclose(res.data, data)
+    res = vis_to_map(vis, shape=(m, n) * apu.pixel, pixel_size=pixel, scheme="natural")
+    # TODO: figure out why atol is needed here
+    assert_allclose(data, res.data, atol=1e-15)
 
-    assert res.reference_coordinate.Tx == pos[0]
-    assert res.reference_coordinate.Ty == pos[1]
-    assert res.scale.axis1 == pixel[0] / apu.pix
-    assert res.scale.axis2 == pixel[1] / apu.pix
+    assert res.reference_coordinate.Tx == pos[1]
+    assert res.reference_coordinate.Ty == pos[0]
+    assert res.scale.axis1 == pixel[1]
+    assert res.scale.axis2 == pixel[0]
     assert res.dimensions.x == m * apu.pix
     assert res.dimensions.y == n * apu.pix
 
 
 def test_to_sunpy_single_pixel_size():
     m = n = 32
-    u = generate_uv(m, pixel_size=2.0 * apu.arcsec)
-    v = generate_uv(m, pixel_size=2.0 * apu.arcsec)
-    u, v = np.meshgrid(u, v)
+    u = generate_uv(m * apu.pix, pixel_size=2.0 * apu.arcsec / apu.pix)
+    v = generate_uv(m * apu.pix, pixel_size=2.0 * apu.arcsec / apu.pix)
+    u, v = np.meshgrid(u, v, indexing="ij")
     uv = np.array([u, v]).reshape(2, m * n) / apu.arcsec
     u, v = uv
 
@@ -232,7 +300,7 @@ def test_to_sunpy_single_pixel_size():
     mp = Map((data, header))
 
     vis = map_to_vis(mp, u=u, v=v)
-    res = vis_to_map(vis, shape=(m, n) * apu.pixel, pixel_size=2 * apu.arcsec)
+    res = vis_to_map(vis, shape=(m, n) * apu.pixel, pixel_size=2 * apu.arcsec / apu.pix)
     assert res.meta["cdelt1"] == 2.0
     assert res.meta["cdelt1"] == 2.0
     assert np.allclose(data, res.data)
@@ -240,9 +308,9 @@ def test_to_sunpy_single_pixel_size():
 
 def test_to_sunpy_map_invalid_pixel_size():
     m = n = 32
-    u = generate_uv(m)
-    v = generate_uv(m)
-    u, v = np.meshgrid(u, v)
+    u = generate_uv(m * apu.pix)
+    v = generate_uv(n * apu.pix)
+    u, v = np.meshgrid(u, v, indexing="ij")
     uv = np.array([u, v]).reshape(2, m * n) / apu.arcsec
     u, v = uv
 
@@ -253,4 +321,4 @@ def test_to_sunpy_map_invalid_pixel_size():
     vis = map_to_vis(mp, u=u, v=v)
 
     with pytest.raises(ValueError):
-        vis_to_map(vis, shape=(m, n) * apu.pixel, pixel_size=[1, 2, 3] * apu.arcsec)
+        vis_to_map(vis, shape=(m, n) * apu.pixel, pixel_size=[1, 2, 3] * apu.arcsec / apu.pix)
diff --git a/xrayvision/tests/test_mem.py b/xrayvision/tests/test_mem.py
index e58471b..f2568a5 100644
--- a/xrayvision/tests/test_mem.py
+++ b/xrayvision/tests/test_mem.py
@@ -1,116 +1,29 @@
 import astropy.units as u
 import numpy as np
 from astropy.convolution.kernels import Gaussian2DKernel
+from numpy.testing import assert_allclose
 
 from xrayvision.imaging import image_to_vis
 from xrayvision.mem import mem, resistant_mean
 
 
 def test_resistant_mean():
+    # fmt: off
     x = np.array(
-        [
-            -0.61566182,
-            -0.3362035,
-            -0.73589487,
-            -0.13014149,
-            -0.91786147,
-            -0.02883366,
-            1.37688696,
-            -0.22389746,
-            -1.37010835,
-            0.21693273,
-            -0.82412723,
-            0.56856134,
-            -0.36824209,
-            -0.65785301,
-            0.10807291,
-            -0.16730037,
-            -0.30308355,
-            0.12605051,
-            -1.14017418,
-            -1.36991364,
-            -0.45275646,
-            0.31636448,
-            -0.1316372,
-            -1.19194891,
-            0.58557167,
-            1.18596057,
-            0.01104923,
-            0.61225975,
-            0.133178,
-            -0.37117463,
-            -0.14882964,
-            -0.59864443,
-            -0.49450362,
-            0.83710754,
-            -0.51270022,
-            0.87380341,
-            0.52904637,
-            0.37973441,
-            0.34061863,
-            -0.10493781,
-            0.86224542,
-            -0.84283578,
-            -0.35069229,
-            -1.29952663,
-            0.39678443,
-            -0.94903521,
-            0.008999,
-            0.52126955,
-            0.24956007,
-            -2.18242896,
-            0.15877853,
-            -0.99771069,
-            -0.14426402,
-            1.48957386,
-            -0.68078661,
-            1.07328579,
-            -2.57370777,
-            2.26526367,
-            1.20964125,
-            -1.26854227,
-            0.20587741,
-            -0.12066582,
-            -0.78971368,
-            -0.36412889,
-            0.12816666,
-            -0.50670234,
-            1.19987195,
-            1.29270178,
-            1.03906824,
-            1.41541003,
-            0.15639539,
-            0.92386204,
-            -2.62109581,
-            0.68909503,
-            -0.34981646,
-            0.93647975,
-            -0.32268089,
-            1.67739013,
-            -0.29216895,
-            -2.30162528,
-            -0.89649313,
-            0.37481482,
-            -2.18661087,
-            0.2190169,
-            -0.30387412,
-            0.84001059,
-            -0.25604226,
-            0.50463104,
-            -1.25674997,
-            0.66314647,
-            -0.0073464,
-            -0.59719724,
-            0.64527612,
-            1.25720287,
-            0.5032624,
-            0.07773715,
-            0.29573575,
-            -0.80567372,
-            -1.92934995,
-            -2.0358119,
-        ]
+        [-0.61566182, -0.3362035, -0.73589487, -0.13014149, -0.91786147, -0.02883366, 1.37688696, -0.22389746,
+         -1.37010835, 0.21693273, -0.82412723, 0.56856134, -0.36824209, -0.65785301, 0.10807291, -0.16730037,
+         -0.30308355, 0.12605051, -1.14017418, -1.36991364, -0.45275646, 0.31636448, -0.1316372, -1.19194891,
+         0.58557167, 1.18596057, 0.01104923, 0.61225975, 0.133178, -0.37117463, -0.14882964, -0.59864443, -0.49450362,
+         0.83710754, -0.51270022, 0.87380341, 0.52904637, 0.37973441, 0.34061863, -0.10493781, 0.86224542, -0.84283578,
+         -0.35069229, -1.29952663, 0.39678443, -0.94903521, 0.008999, 0.52126955, 0.24956007, -2.18242896, 0.15877853,
+         -0.99771069, -0.14426402, 1.48957386, -0.68078661, 1.07328579, -2.57370777, 2.26526367, 1.20964125,
+         -1.26854227, 0.20587741, -0.12066582, -0.78971368, -0.36412889, 0.12816666, -0.50670234, 1.19987195,
+         1.29270178, 1.03906824, 1.41541003, 0.15639539, 0.92386204, -2.62109581, 0.68909503, -0.34981646, 0.93647975,
+         -0.32268089, 1.67739013, -0.29216895, -2.30162528, -0.89649313, 0.37481482, -2.18661087, 0.2190169,
+         -0.30387412, 0.84001059, -0.25604226, 0.50463104, -1.25674997, 0.66314647, -0.0073464, -0.59719724, 0.64527612,
+         1.25720287, 0.5032624, 0.07773715, 0.29573575, -0.80567372, -1.92934995, -2.0358119,]
     )
+    # fmt: on
     rmean, rsigma = resistant_mean(x, 3)
     # values for IDL routine "RESISTANT_Mean, xx, 3, rmean, rsig, num"
     assert np.allclose(rmean, -0.10947954)
@@ -143,9 +56,11 @@ def test_mem():
     sub_uv = np.vstack([x.flatten(), y.flatten()]) / u.arcsec
     # sub_uv = np.hstack([sub_uv, np.zeros((2, 1))]) / u.arcsec
 
-    vis = image_to_vis(data, u=sub_uv[0, :], v=sub_uv[1, :], pixel_size=2 * u.arcsec)
+    vis = image_to_vis(data * u.dimensionless_unscaled, u=sub_uv[0, :], v=sub_uv[1, :], pixel_size=2 * u.arcsec / u.pix)
     setattr(vis, "amplitude_error", np.sqrt(np.abs(vis.vis)))
     setattr(vis, "label", [str(x) for x in np.sqrt(x**2 + y**2).flatten()])
 
-    res = mem(vis, percent_lambda=None, shape=(m, n) * u.pix, pixel=[2, 2] * u.arcsec, maxiter=1000, tol=1e-3)
-    np.testing.assert_allclose(res.data - data / 4, 0.0, atol=1e-3)
+    res = mem(
+        vis, percent_lambda=None, shape=(m, n) * u.pix, pixel_size=[2, 2] * u.arcsec / u.pix, maxiter=1000, tol=1e-3
+    )
+    assert_allclose(res.data, data, atol=1e-1)
diff --git a/xrayvision/tests/test_transform.py b/xrayvision/tests/test_transform.py
index 97fe9c9..4fcb6d8 100644
--- a/xrayvision/tests/test_transform.py
+++ b/xrayvision/tests/test_transform.py
@@ -1,8 +1,8 @@
 import astropy.units as apu
 import numpy as np
 import pytest
-from astropy.convolution import Gaussian2DKernel
-from numpy.testing import assert_allclose, assert_array_equal
+from numpy.fft import fft2, fftshift, ifft2, ifftshift
+from numpy.testing import assert_allclose
 from scipy import signal
 
 from xrayvision.transform import dft_map, generate_uv, generate_xy, idft_map
@@ -10,268 +10,295 @@
 
 @pytest.mark.parametrize("pixel_size", [0.5, 1, 2, 3])
 def test_generate_xy_pixel_size(pixel_size):
-    # Odd
-    pixel_size = pixel_size * apu.arcsec
-    odd = np.linspace(-16 * pixel_size, 16 * pixel_size, 33)
-    assert np.array_equal(odd, generate_xy(33, pixel_size=pixel_size))
+    odd = np.linspace(-16 * pixel_size, 16 * pixel_size, 33) * apu.arcsec
+    even = np.linspace(-15.5 * pixel_size, 15.5 * pixel_size, 32) * apu.arcsec
+    pixel_size = pixel_size * apu.arcsec / apu.pix
 
-    # Even
-    even = np.linspace(-15.5 * pixel_size, 15.5 * pixel_size, 32)
-    assert np.array_equal(even, generate_xy(32, pixel_size=pixel_size))
+    assert np.array_equal(odd, generate_xy(33 * apu.pix, pixel_size=pixel_size))
+    assert np.array_equal(even, generate_xy(32 * apu.pix, pixel_size=pixel_size))
 
 
-@pytest.mark.parametrize("center", [0, +5.5, -5.5])
-def test_generate_xy_offset(center):
-    # Odd
-    center = center * apu.arcsec
-    odd = np.linspace(-16, 16, 33) * apu.arcsec + center
-    assert np.array_equal(odd, generate_xy(33, center=center))
+@pytest.mark.parametrize("phase_centre", [0, +5.5, -5.5])
+def test_generate_xy_offset(phase_centre):
+    phase_centre = phase_centre * apu.arcsec
+    even = np.linspace(-15.5, 15.5, 32) * apu.arcsec + phase_centre
+    odd = np.linspace(-16, 16, 33) * apu.arcsec + phase_centre
 
-    # Even
-    even = np.linspace(-15.5, 15.5, 32) * apu.arcsec + center
-    assert np.array_equal(even, generate_xy(32, center=center))
+    assert np.array_equal(even, generate_xy(32 * apu.pix, phase_centre=phase_centre))
+    assert np.array_equal(odd, generate_xy(33 * apu.pix, phase_centre=phase_centre))
 
 
-@pytest.mark.parametrize("center, pixel_size", [(0, (0.5, 1, 2, 3)), (+5.5, (0.5, 1, 2, 3)), (-5.5, (0.5, 1, 2, 3))])
-def test_generate_xy_offset_size(center, pixel_size):
-    center = center * apu.arcsec
-    pixel_size = pixel_size * apu.arcsec
+@pytest.mark.parametrize(
+    "phase_centre, pixel_size", [(0, (0.5, 1, 2, 3)), (+5.5, (0.5, 1, 2, 3)), (-5.5, (0.5, 1, 2, 3))]
+)
+def test_generate_xy_offset_size(phase_centre, pixel_size):
+    phase_centre = phase_centre * apu.arcsec
+    pixel_size = pixel_size * apu.arcsec / apu.pix
     # No sure if this is a good idea could be hard to debug
     for size in pixel_size:
         # Odd
-        odd = np.linspace(-16 * size, 16 * size, 33) + center
-        assert np.array_equal(odd, generate_xy(33, center=center, pixel_size=size))
+        odd = np.linspace(-16 * size, 16 * size, 33) * apu.pix + phase_centre
+        assert np.array_equal(odd, generate_xy(33 * apu.pix, phase_centre=phase_centre, pixel_size=size))
 
         # Even
-        even = np.linspace(-15.5 * size, 15.5 * size, 32) + center
-        assert np.array_equal(even, generate_xy(32, center=center, pixel_size=size))
+        even = np.linspace(-15.5 * size, 15.5 * size, 32) * apu.pix + phase_centre
+        assert np.array_equal(even, generate_xy(32 * apu.pix, phase_centre=phase_centre, pixel_size=size))
 
 
 @pytest.mark.parametrize("pixel_size", [0.5, 1, 2, 3])
 def test_generate_uv_pixel_size(pixel_size):
-    pixel_size = pixel_size * apu.arcsec
+    pixel_size = pixel_size * apu.arcsec / apu.pix
     m = 33
     n = 32
 
     # Odd
-    odd = np.linspace(-((m - 1) / 2) * (1 / (m * pixel_size)), ((m - 1) / 2) * (1 / (m * pixel_size)), m)
+    odd = np.linspace(-((m - 1) / 2) * (1 / (m * pixel_size)), ((m - 1) / 2) * (1 / (m * pixel_size)), m) / apu.pix
 
     # Know issue with quantities and isclose/allclose need to add unit to atol default value
-    assert np.allclose(odd, generate_uv(m, pixel_size=pixel_size), atol=1e-8 / apu.arcsec)
+    assert np.allclose(odd, generate_uv(m * apu.pix, pixel_size=pixel_size), atol=1e-8 / apu.arcsec)
 
     # Even
-    even = (np.arange(n) - n / 2 + 0.5) * (1 / (pixel_size * n))
-    assert np.allclose(even, generate_uv(32, pixel_size=pixel_size), atol=1e-8 / apu.arcsec)
+    even = (np.arange(n) - n / 2 + 0.5) * (1 / (pixel_size * n)) / apu.pix
+    assert np.allclose(even, generate_uv(32 * apu.pix, pixel_size=pixel_size), atol=1e-8 / apu.arcsec)
 
 
-@pytest.mark.parametrize("center", [0.0, -5.5, 5.5])
-def test_generate_uv_pixel_offset(center):
-    center = center * apu.arcsec
+@pytest.mark.parametrize("phase_centre", [0.0, -5.5, 5.5])
+def test_generate_uv_pixel_offset(phase_centre):
+    phase_centre = phase_centre * apu.arcsec
     m = 33
     n = 32
 
     # Odd
     odd = np.linspace(-((m - 1) / 2) * (1 / m), ((m - 1) / 2) * (1 / m), m) / apu.arcsec
-    if center.value != 0:
-        odd += 1 / center
-    assert np.allclose(odd, generate_uv(m, center=center), atol=1e-8 / apu.arcsec)
+    if phase_centre.value != 0:
+        odd += 1 / phase_centre
+    assert np.allclose(odd, generate_uv(m * apu.pix, phase_centre=phase_centre), atol=1e-8 / apu.arcsec)
 
     # Even
     even = (np.arange(n) - n / 2 + 0.5) * (1 / n) / apu.arcsec
-    if center.value != 0:
-        even += 1 / center
-    assert np.allclose(even, generate_uv(32, center=center), atol=1e-8 / apu.arcsec)
+    if phase_centre.value != 0:
+        even += 1 / phase_centre
+    assert np.allclose(even, generate_uv(n * apu.pix, phase_centre=phase_centre), atol=1e-8 / apu.arcsec)
 
 
-@pytest.mark.parametrize("center, pixel_size", [(0, (0.5, 1, 2, 3)), (+5.5, (0.5, 1, 2, 3)), (-5.5, (0.5, 1, 2, 3))])
-def test_generate_uv_offset_size(center, pixel_size):
-    center = center * apu.arcsec
-    pixel_size = pixel_size * apu.arcsec
+@pytest.mark.parametrize(
+    "phase_centre, pixel_size", [(0, (0.5, 1, 2, 3)), (+5.5, (0.5, 1, 2, 3)), (-5.5, (0.5, 1, 2, 3))]
+)
+def test_generate_uv_offset_size(phase_centre, pixel_size):
+    phase_centre = phase_centre * apu.arcsec
+    pixel_size = pixel_size * apu.arcsec / apu.pix
     m = 33
     n = 32
 
     for size in pixel_size:
         # Odd
-        odd = np.linspace(-((m - 1) / 2) * (1 / (size * m)), ((m - 1) / 2) * (1 / (size * m)), m)
-        if center != 0:
-            odd += 1 / center
-        assert np.allclose(odd, generate_uv(m, center=center, pixel_size=size), atol=1e-8 / apu.arcsec)
+        odd = np.linspace(-((m - 1) / 2) * (1 / (size * m)), ((m - 1) / 2) * (1 / (size * m)), m) / apu.pix
+        if phase_centre != 0:
+            odd += 1 / phase_centre
+        assert np.allclose(
+            odd, generate_uv(m * apu.pix, phase_centre=phase_centre, pixel_size=size), atol=1e-8 / apu.arcsec
+        )
 
         # Even
-        even = (np.arange(n) - n / 2 + 0.5) * (1 / (size * n))
-        if center != 0:
-            even += 1 / center
-        assert np.allclose(even, generate_uv(32, center=center, pixel_size=size), atol=1e-8 / apu.arcsec)
-
-
-@pytest.mark.parametrize("shape", [(33, 33), (32, 32), (33, 32), (32, 33)])
-def test_dft_idft_map(shape):
-    m, n = shape
-    m * n
-    uu = generate_uv(n)
-    vv = generate_uv(m)
-    u, v = np.meshgrid(uu, vv)
+        even = (np.arange(n) - n / 2 + 0.5) * (1 / (size * n)) / apu.pix
+        if phase_centre != 0:
+            even += 1 / phase_centre
+        assert np.allclose(
+            even, generate_uv(n * apu.pix, phase_centre=phase_centre, pixel_size=size), atol=1e-8 / apu.arcsec
+        )
+
+
+@pytest.mark.parametrize("shape", [(3, 3), (2, 3), (3, 2), (2, 2)])
+@pytest.mark.parametrize("pixel_size", [(1, 1), (1, 2), (2, 1), (2, 2)])
+@pytest.mark.parametrize("center", [(0, 0), (-1, 1), (1, -1), (1, 1)])
+def test_dft_idft(shape, pixel_size, center):
+    data = np.arange(np.prod(shape)).reshape(shape)
+    uu = generate_uv(
+        shape[0] * apu.pix, phase_centre=center[0] * apu.arcsec, pixel_size=pixel_size[0] * apu.arcsec / apu.pix
+    )
+    vv = generate_uv(
+        shape[1] * apu.pix, phase_centre=center[1] * apu.arcsec, pixel_size=pixel_size[1] * apu.arcsec / apu.pix
+    )
+    u, v = np.meshgrid(uu, vv, indexing="ij")
     u = u.flatten()
     v = v.flatten()
 
-    # All zeros
-    zeros = np.zeros((m, n))
-    vis = dft_map(zeros, u=u, v=v)
-    # All visibilities should be zero
-    # assert_array_equal(np.zeros(np.prod((m, n)), dtype=complex), vis)
-    out_map = idft_map(vis, u=u, v=v, shape=shape)
-    # Should get back original map after dft(idft())
-    assert_allclose(zeros, out_map / np.prod((m, n)))
-
-    # All ones
-    ones = np.ones((m, n))
-    vis = dft_map(ones, u=u, v=v)
-    # All visibilities should be 1
-    # assert_allclose(np.ones(np.prod((m, n)), dtype=complex), vis)
-    out_map = idft_map(vis, u=u, v=v, shape=(m, n))
-    assert_allclose(ones, out_map / np.prod((m, n)))
-
-    # Delta
-    delta = zeros[:, :]
-    delta[m // 2, n // 2] = 1.0
-    vis = dft_map(delta, u=u, v=v)
-    out_map = idft_map(vis, u=u, v=v, shape=shape)
-    assert_allclose(out_map / np.prod(shape), delta, atol=1e-14)
-
-    # Gaussian - astropy has axis in reverse order compared to numpy
-    gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array
-    vis = dft_map(gaussian, u=u, v=v)
-    out_map = idft_map(vis, u=u, v=v, shape=shape)
-    assert_allclose(out_map / np.prod(shape), gaussian)
-
-
-# @pytest.mark.skip('UV coordinate generation off somewhere')
-@pytest.mark.parametrize("pixel_size", [(1.0, 1.0), (0.5, 0.5), (1.5, 1.5), (2.0, 0.5), (0.5, 2.0)])
-def test_dft_idft_map_pixel_size(pixel_size):
-    pixel_size = pixel_size * apu.arcsec
-    shape = (32, 32)
-    m, n = shape
-    size = m * n
-    uu = generate_uv(m, pixel_size=pixel_size[0])
-    vv = generate_uv(n, pixel_size=pixel_size[1])
-    u, v = np.meshgrid(uu, vv)
-    u = u.flatten()
-    v = v.flatten()
-
-    # All zeros
-    zeros = np.zeros(shape)
-    vis = dft_map(zeros, u=u, v=v, pixel_size=pixel_size)
-    # All visibilities should be zero
-    assert_array_equal(np.zeros(size, dtype=complex), vis)
-    out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
-    # Should get back original map after dft(idft())
-    assert_array_equal(out_map / np.prod(shape), zeros)
-
-    # All ones
-    ones = np.ones(shape)
-    vis = dft_map(ones, u=u, v=v, pixel_size=pixel_size)
-    out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
-    assert_allclose(out_map / np.prod(shape), ones)
-
-    # Delta
-    delta = zeros[:, :]
-    delta[m // 2, n // 2] = 1.0
-    vis = dft_map(delta, u=u, v=v, pixel_size=pixel_size)
-    out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
-    assert_allclose(out_map / np.prod(shape), delta, atol=1e-14)
-
-    # Gaussian - astropy has axis in reverse order compared to numpy
-    gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array
-    vis = dft_map(gaussian, u=u, v=v, pixel_size=pixel_size)
-    out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
-    assert_allclose(out_map / np.prod(shape), gaussian)
-
-
-# @pytest.mark.skip('UV coordinate generation off somewhere')
-@pytest.mark.parametrize("center", [(0, 0), (2.1, 2.1), (5.4, -5.4), (-5.6, 5.6)])
-def test_dft_idft_map_center(center):
-    center = center * apu.arcsec
-    shape = (33, 33)
-    m, n = shape
-    size = m * n
-    uu = generate_uv(n, center=center[1])
-    vv = generate_uv(m, center=center[0])
-    u, v = np.meshgrid(uu, vv)
-    u = u.flatten()
-    v = v.flatten()
-
-    # All zeros
-    zeros = np.zeros(shape)
-    vis = dft_map(zeros, u=u, v=v, center=center)
-    # All visibilities should be zero
-    assert_array_equal(vis, np.zeros(size, dtype=complex))
-    out_map = idft_map(vis, u=u, v=v, shape=shape, center=center)
-    # Should get back original map after dft(idft())
-    assert_array_equal(out_map / np.prod(shape), zeros)
-
-    # All ones
-    ones = np.ones(shape)
-    vis = dft_map(ones, u=u, v=v, center=center)
-    out_map = idft_map(vis, u=u, v=v, shape=shape, center=center)
-    assert_allclose(out_map / np.prod(shape), ones)
-
-    # Delta
-    delta = zeros[:, :]
-    delta[m // 2, n // 2] = 1.0
-    vis = dft_map(delta, u=u, v=v, center=center)
-    out_map = idft_map(vis, u=u, v=v, shape=shape, center=center)
-    assert_allclose(out_map / np.prod(shape), delta, atol=1e-14)
-
-    # Gaussian - astropy has axis in reverse order compared to numpy
-    gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array
-    vis = dft_map(gaussian, u=u, v=v, center=center)
-    out_map = idft_map(vis, u=u, v=v, shape=shape, center=center)
-    assert_allclose(out_map / np.prod(shape), gaussian)
-
-
-# @pytest.mark.skip('UV coordinate generation off somewhere')
-@pytest.mark.parametrize(
-    "shape, pixel_size",
-    [((11, 10), (0.5, 0.5)), ((11, 10), (1.0, 1.0)), ((11, 10), (2.0, 2.0)), ((11, 10), (3.0, 3.0))],
-)
-def test_dft_idft_map_shape_pixel_size(shape, pixel_size):
-    pixel_size = pixel_size * apu.arcsec
-    m, n = shape
-    size = m * n
-
-    vv = generate_uv(m, pixel_size=pixel_size[0])
-    uu = generate_uv(n, pixel_size=pixel_size[1])
-    u, v = np.meshgrid(uu, vv)
-    u = u.flatten()
-    v = v.flatten()
-
-    # All zeros
-    zeros = np.zeros(shape)
-    vis = dft_map(zeros, u=u, v=v, pixel_size=pixel_size)
-    # All visibilities should be zero
-    assert_array_equal(np.zeros(size, dtype=complex), vis)
-    out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
-    # Should get back original map after dft(idft())
-    assert_array_equal(out_map / np.prod(shape), zeros)
-
-    # All ones
-    ones = np.ones(shape)
-    vis = dft_map(ones, u=u, v=v, pixel_size=pixel_size)
-    out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
-    assert_allclose(out_map / np.prod(shape), ones)
-
-    # Delta
-    delta = zeros[:, :]
-    delta[m // 2, n // 2] = 1.0
-    vis = dft_map(delta, u=u, v=v, pixel_size=pixel_size)
-    out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
-    assert_allclose(out_map / np.prod(shape), delta, atol=1e-14)
-
-    # Gaussian - astropy has axis in reverse order compared to numpy
-    gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array
-    vis = dft_map(gaussian, u=u, v=v, pixel_size=pixel_size)
-    out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
-    assert_allclose(out_map / np.prod(shape), gaussian)
+    vis = dft_map(data, u=u, v=v, pixel_size=pixel_size * apu.arcsec / apu.pix)
+    img = idft_map(
+        vis, u=u, v=v, weights=1 / np.prod(shape), shape=shape * apu.pix, pixel_size=pixel_size * apu.arcsec / apu.pix
+    )
+    assert_allclose(img, data, atol=1e-13)
+
+
+# @pytest.mark.parametrize("shape", [(33, 33), (32, 32), (33, 32), (32, 33)])
+# def test_dft_idft_map(shape):
+#     m, n = shape
+#     uu = generate_uv(m * apu.pix)
+#     vv = generate_uv(n * apu.pix)
+#     u, v = np.meshgrid(uu, vv, indexing="ij")
+#     u = u.flatten()
+#     v = v.flatten()
+#
+#     # All zeros
+#     zeros = np.zeros((m, n))
+#     vis = dft_map(zeros, u=u, v=v)
+#     # All visibilities should be zero
+#     # assert_array_equal(np.zeros(np.prod((m, n)), dtype=complex), vis)
+#     out_map = idft_map(vis, u=u, v=v, shape=shape * apu.pix)
+#     # Should get back the original map after dft(idft())
+#     assert_allclose(zeros, out_map / np.prod((m, n)))
+#
+#     # All ones
+#     ones = np.ones((m, n))
+#     vis = dft_map(ones, u=u, v=v)
+#     # All visibilities should be 1
+#     # assert_allclose(np.ones(np.prod((m, n)), dtype=complex), vis)
+#     out_map = idft_map(vis, u=u, v=v, shape=shape * apu.pix)
+#     assert_allclose(ones, out_map / np.prod((m, n)))
+#
+#     # Delta
+#     delta = zeros[:, :]
+#     delta[m // 2, n // 2] = 1.0
+#     vis = dft_map(delta, u=u, v=v)
+#     out_map = idft_map(vis, u=u, v=v, shape=shape * apu.pix)
+#     assert_allclose(out_map / np.prod(shape), delta, atol=1e-14)
+#
+#     # Gaussian - astropy has axis in reverse order compared to numpy
+#     gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array
+#     vis = dft_map(gaussian, u=u, v=v)
+#     out_map = idft_map(vis, u=u, v=v, shape=shape * apu.pix)
+#     assert_allclose(out_map / np.prod(shape), gaussian)
+#
+#
+# # @pytest.mark.skip('UV coordinate generation off somewhere')
+# @pytest.mark.parametrize("pixel_size", [(1.0, 1.0), (0.5, 0.5), (1.5, 1.5), (2.0, 0.5), (0.5, 2.0)])
+# def test_dft_idft_map_pixel_size(pixel_size):
+#     pixel_size = pixel_size * apu.arcsec / apu.pix
+#     shape = (32, 32) * apu.pix
+#     m, n = shape.value.astype(int)
+#     size = m * n
+#     uu = generate_uv(m * apu.pix, pixel_size=pixel_size[0])
+#     vv = generate_uv(n * apu.pix, pixel_size=pixel_size[1])
+#     u, v = np.meshgrid(uu, vv)
+#     u = u.flatten()
+#     v = v.flatten()
+#
+#     # All zeros
+#     zeros = np.zeros(shape.value.astype(int))
+#     vis = dft_map(zeros, u=u, v=v, pixel_size=pixel_size)
+#     # All visibilities should be zero
+#     assert_array_equal(np.zeros(size, dtype=complex), vis)
+#     out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
+#     # Should get back the original map after dft(idft())
+#     assert_array_equal(out_map / np.prod(shape.value), zeros)
+#
+#     # All ones
+#     ones = np.ones(shape.value.astype(int))
+#     vis = dft_map(ones, u=u, v=v, pixel_size=pixel_size)
+#     out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
+#     assert_allclose(out_map / np.prod(shape.value), ones)
+#
+#     # Delta
+#     delta = zeros[:, :]
+#     delta[m // 2, n // 2] = 1.0
+#     vis = dft_map(delta, u=u, v=v, pixel_size=pixel_size)
+#     out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
+#     assert_allclose(out_map / np.prod(shape.value), delta, atol=1e-14)
+#
+#     # Gaussian - astropy has axis in reverse order compared to numpy
+#     gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array
+#     vis = dft_map(gaussian, u=u, v=v, pixel_size=pixel_size)
+#     out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
+#     assert_allclose(out_map / np.prod(shape.value), gaussian)
+#
+#
+# # @pytest.mark.skip('UV coordinate generation off somewhere')
+# @pytest.mark.parametrize("phase_centre", [(0, 0), (2.1, 1.1), (5.4, -4.5), (-5.6, 5.6)])
+# def test_dft_idft_map_center(phase_centre):
+#     phase_centre = phase_centre * apu.arcsec
+#     shape = (33, 33)
+#     m, n = shape
+#     size = m * n
+#     shape = shape * apu.pix
+#     uu = generate_uv(n * apu.pix, phase_centre=phase_centre[0])
+#     vv = generate_uv(m * apu.pix, phase_centre=phase_centre[1])
+#     u, v = np.meshgrid(uu, vv)
+#     u = u.flatten()
+#     v = v.flatten()
+#
+#     # All zeros
+#     zeros = np.zeros((m, n))
+#     vis = dft_map(zeros, u=u, v=v, phase_centre=phase_centre)
+#     # All visibilities should be zero
+#     assert_array_equal(vis, np.zeros(size, dtype=complex))
+#     out_map = idft_map(vis, u=u, v=v, shape=shape, phase_centre=phase_centre)
+#     # Should get back the original map after dft(idft())
+#     assert_array_equal(out_map / np.prod((m, n)), zeros)
+#
+#     # All ones
+#     ones = np.ones((m, n))
+#     vis = dft_map(ones, u=u, v=v, phase_centre=phase_centre)
+#     out_map = idft_map(vis, u=u, v=v, shape=shape, phase_centre=phase_centre)
+#     assert_allclose(out_map / np.prod((m, n)), ones)
+#
+#     # Delta
+#     delta = zeros[:, :]
+#     delta[m // 2, n // 2] = 1.0
+#     vis = dft_map(delta, u=u, v=v, phase_centre=phase_centre)
+#     out_map = idft_map(vis, u=u, v=v, shape=shape, phase_centre=phase_centre)
+#     assert_allclose(out_map / np.prod((m, n)), delta, atol=1e-14)
+#
+#     # Gaussian - astropy has axis in reverse order compared to numpy
+#     gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array
+#     vis = dft_map(gaussian, u=u, v=v, phase_centre=phase_centre)
+#     out_map = idft_map(vis, u=u, v=v, shape=shape, phase_centre=phase_centre)
+#     assert_allclose(out_map / np.prod((m, n)), gaussian)
+#
+#
+# # @pytest.mark.skip('UV coordinate generation off somewhere')
+# @pytest.mark.parametrize(
+#     "shape, pixel_size",
+#     [((11, 10), (0.5, 0.5)), ((11, 10), (1.0, 1.0)), ((11, 10), (2.0, 2.0)), ((11, 10), (3.0, 3.0))],
+# )
+# def test_dft_idft_map_shape_pixel_size(shape, pixel_size):
+#     pixel_size = pixel_size * apu.arcsec / apu.pix
+#     m, n = shape
+#     size = m * n
+#     shape = shape * apu.pix
+#
+#     uu = generate_uv(m * apu.pix, pixel_size=pixel_size[0])
+#     vv = generate_uv(n * apu.pix, pixel_size=pixel_size[1])
+#     u, v = np.meshgrid(uu, vv, indexing="ij")
+#     u = u.flatten()
+#     v = v.flatten()
+#
+#     # All zeros
+#     zeros = np.zeros((m, n))
+#     vis = dft_map(zeros, u=u, v=v, pixel_size=pixel_size)
+#     # All visibilities should be zero
+#     assert_array_equal(np.zeros(size, dtype=complex), vis)
+#     out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
+#     # Should get back the original map after dft(idft())
+#     assert_array_equal(out_map / np.prod((m, n)), zeros)
+#
+#     # All ones
+#     ones = np.ones((m, n))
+#     vis = dft_map(ones, u=u, v=v, pixel_size=pixel_size)
+#     out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
+#     assert_allclose(out_map / np.prod((m, n)), ones)
+#
+#     # Delta
+#     delta = zeros[:, :]
+#     delta[m // 2, n // 2] = 1.0
+#     vis = dft_map(delta, u=u, v=v, pixel_size=pixel_size)
+#     out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
+#     assert_allclose(out_map / np.prod((m, n)), delta, atol=1e-14)
+#
+#     # Gaussian - astropy has axis in reverse order compared to numpy
+#     gaussian = Gaussian2DKernel(5, x_size=n, y_size=m).array
+#     vis = dft_map(gaussian, u=u, v=v, pixel_size=pixel_size)
+#     out_map = idft_map(vis, u=u, v=v, shape=shape, pixel_size=pixel_size)
+#     assert_allclose(out_map / np.prod((m, n)), gaussian)
 
 
 def test_equivalence_of_convolve():
@@ -280,10 +307,9 @@ def test_equivalence_of_convolve():
     data[3:6, 3:6] = 5.0
 
     m, n = (33, 33)
-    m * n
 
-    vv = generate_uv(m)
-    uu = generate_uv(n)
+    vv = generate_uv(m * apu.pix)
+    uu = generate_uv(n * apu.pix)
     u, v = np.meshgrid(uu, vv)
     u = u.flatten()
     v = v.flatten()
@@ -298,18 +324,87 @@ def test_equivalence_of_convolve():
 
     # bp1 = idft_map(full_vis, u=u, v=v, shape=(33, 33))
 
-    bp2 = idft_map(sub_vis[non_zero], u=u[non_zero], v=v[non_zero], shape=(33, 33))
+    bp2 = idft_map(sub_vis[non_zero], u=u[non_zero], v=v[non_zero], shape=(33, 33) * apu.pix)
 
-    # Need to make the psf large enough to slide over entire data window
-    psf1 = idft_map(sampling[non_zero], u=u[non_zero], v=v[non_zero], shape=(33 * 3, 33 * 3))
+    # Need to make the psf is large enough to slide over the entire data window
+    psf1 = idft_map(sampling[non_zero], u=u[non_zero], v=v[non_zero], shape=(33 * 3, 33 * 3) * apu.pix)
 
     conv = signal.convolve(data, psf1, mode="same", method="fft")
 
-    psf2 = idft_map(sampling[non_zero], u=u[non_zero], v=v[non_zero], shape=(33, 33))
+    psf2 = idft_map(sampling[non_zero], u=u[non_zero], v=v[non_zero], shape=(33, 33) * apu.pix)
 
-    bp3 = idft_map(full_vis[non_zero], u=u[non_zero], v=v[non_zero], shape=(33, 33))
+    bp3 = idft_map(full_vis[non_zero], u=u[non_zero], v=v[non_zero], shape=(33, 33) * apu.pix)
 
     assert np.allclose(bp2, conv)
     assert np.allclose(bp2, bp3)
-    # Due to the enlarged psf need to only use center portion
+    # Due to the enlarged psf need to only use the centre portion
     assert np.allclose(psf1[33:66, 33:66], psf2)
+
+
+def test_phase_centre_equivalence():
+    # Calculate full u, v coverage so will be equivalent to a discrete Fourier transform (DFT)
+    data = np.random.randn(8, 8) / apu.arcsec**2
+    u = generate_uv(8 * apu.pix, phase_centre=0 * apu.arcsec, pixel_size=1 * apu.arcsec / apu.pix)
+    v = generate_uv(8 * apu.pix, phase_centre=0 * apu.arcsec, pixel_size=1 * apu.arcsec / apu.pix)
+    u, v = np.meshgrid(u, v)
+    u, v = np.array([u, v]).reshape(2, u.size) / apu.arcsec
+
+    # calculate normal 0,0
+    vis = dft_map(data, u=u, v=v, pixel_size=[1, 1] * apu.arcsec / apu.pix)
+    img1 = idft_map(
+        vis, u=u, v=v, weights=1 / u.size, pixel_size=[1, 1] * apu.arcsec / apu.pix, shape=data.shape * apu.pix
+    )
+    assert_allclose(data, img1)
+
+    # extract phase and amp
+    phase = np.arctan2(np.imag(vis), np.real(vis))
+    amp = np.abs(vis)
+
+    # change vis to a phase centre of 5, 5
+    phase_shift = 2 * np.pi * (5 * apu.arcsec * u + 5 * apu.arcsec * v) * apu.rad
+    vis_shifted = (np.cos(phase + phase_shift) + np.sin(phase + phase_shift) * 1j) * amp
+
+    # make image with centre of 5, 5 with shifted vis
+    img2 = idft_map(
+        vis_shifted,
+        u=u,
+        v=v,
+        weights=1 / u.size,
+        pixel_size=[1, 1] * apu.arcsec / apu.pix,
+        shape=data.shape * apu.pix,
+        phase_centre=[5, 5] * apu.arcsec,
+    )
+    assert np.allclose(data, img2)
+
+
+def test_fft_equivalence():
+    # Odd (3, 3) so symmetric and chose shape and pixel size so xy values will run from 0 to 2 the same as in fft
+    # TODO: add same kind of test for even for fft2 then A[n/2] has both pos and negative nyquist frequencies
+    #  e.g shape (2, 2), (3, 2), (2, 3)
+    shape = (3, 3)
+    pixel = (1, 1)
+    center = (1, 1)
+
+    data = np.arange(np.prod(shape)).reshape(shape)
+    uu = generate_uv(
+        shape[0] * apu.pix, phase_centre=center[0] * apu.arcsec, pixel_size=pixel[0] * apu.arcsec / apu.pix
+    )
+    vv = generate_uv(
+        shape[1] * apu.pix, phase_centre=center[1] * apu.arcsec, pixel_size=pixel[1] * apu.arcsec / apu.pix
+    )
+    u, v = np.meshgrid(uu, vv, indexing="ij")
+    u = u.flatten()
+    v = v.flatten()
+
+    vis = dft_map(data, u=u, v=v, pixel_size=pixel * apu.arcsec / apu.pix, phase_centre=center * apu.arcsec)
+
+    ft = fft2(data)
+    fts = fftshift(ft)
+    vis = vis.reshape(shape)
+    # Convention in xrayvison has the minus sign on the forward transform but numpy has it on reverse
+    vis_conj = np.conjugate(vis)
+    assert_allclose(fts, vis_conj, atol=1e-13)
+
+    vis_ft = ifftshift(vis_conj)
+    img = ifft2(vis_ft)
+    assert_allclose(np.real(img), data, atol=1e-14)
diff --git a/xrayvision/tests/test_visibility.py b/xrayvision/tests/test_visibility.py
index b47ab04..cfa1d0e 100644
--- a/xrayvision/tests/test_visibility.py
+++ b/xrayvision/tests/test_visibility.py
@@ -1,16 +1,19 @@
-from pathlib import Path
-
 import astropy.units as apu
 import pytest
 from astropy.tests.helper import assert_quantity_allclose
+from numpy.testing import assert_array_equal
 
 import xrayvision.visibility as vm
+from xrayvision.visibility import Visibility
 
 
-@pytest.fixture
-def test_data_dir():
-    path = Path(__file__).parent.parent / "data"
-    return path
+def test_visibility():
+    vis = Visibility(vis=1 * apu.ct, u=1 / apu.deg, v=1 / apu.deg)
+    assert vis.vis == 1 * apu.ct
+    assert vis.u == 1 / apu.deg
+    assert vis.v == 1 / apu.deg
+    assert_array_equal([0, 0] * apu.arcsec, vis.phase_centre)
+    assert_array_equal([0, 0] * apu.arcsec, vis.offset)
 
 
 @pytest.fixture