Repository: LouisDesdoigts/dLux (language: python)

Issue:
Write general docs



Discussion / hints:

Fix patch:
diff --git a/src/__init__.py b/src/__init__.py
index 82dad8db..537e1d81 100644
--- a/src/__init__.py
+++ b/src/__init__.py
@@ -1,7 +1,11 @@
 name = "dLux"
 __version__ = "0.1.2"
+__author__ = "Louis Desdoigts"
+__author__ = "Jordan Dennis"
+__author__ = "Benjamin Pope"
 
 from .base import *
 from .layers import *
 from .detector_layers import *
 from .propagators import *
+from .wavefronts import *
diff --git a/src/base.py b/src/base.py
index cb7841e1..86c13e27 100644
--- a/src/base.py
+++ b/src/base.py
@@ -4,175 +4,7 @@
 import equinox as eqx
 from copy import deepcopy
 
-class PhysicalWavefront(eqx.Module):
-    """
-    Class to store the optical information and paramaters
-    """
-    amplitude:  np.ndarray
-    phase:      np.ndarray
-    wavel:      float
-    offset:     np.ndarray
-    pixelscale: float
-    planetype:  str
-    
-    def __init__(self, wavel, offset):
-        self.wavel = np.array(wavel).astype(float)
-        self.offset = np.array(offset).astype(float)
-        self.planetype = "Pupil"
-        self.pixelscale = None
-        self.amplitude = None
-        self.phase = None
-        
-        
-        
-        
-    """Real and Imag"""
-    def get_real(self):
-        return self.ampltidue * np.cos(self.phase)
-    
-    def get_imag(self):
-        return self.ampltidue * np.sin(self.phase)
-    
-    
-    
-    
-    """Updating the Phasor"""
-    def multiply_ampl(self, array_like):
-        new_ampl = self.amplitude * array_like
-        return eqx.tree_at(lambda wf: wf.amplitude, self, new_ampl)
-        
-    def add_phase(self, array_like):
-        new_phase = self.phase + array_like
-        return eqx.tree_at(lambda wf: wf.phase, self, new_phase)
-        
-    def update_phasor(self, amplitude, phase):
-        self = eqx.tree_at(lambda wf: wf.amplitude, self, amplitude)
-        self = eqx.tree_at(lambda wf: wf.phase,     self, phase)
-        return self
-    
-    
-    
-    """Optics stuff"""
-    def wf2psf(self):
-        return self.amplitude**2
-        
-    def add_opd(self, opd):
-        phase = 2*np.pi*opd/self.wavel
-        return self.add_phase(phase)
-        
-    def normalise(self):
-        total_intensity = np.sqrt(np.sum(self.amplitude**2))
-        return self.multiply_ampl(1/total_intensity)
-    
-    
-    
-    
-        
-    """Coordinates"""    
-    def get_xs_vec(self, npix):
-        """ Returns a paraxial 1d vector """
-        return np.arange(npix) - (npix-1)/2 
-    
-    def get_XXYY(self):
-        """ """
-        xs = self.get_xs_vec(self.amplitude.shape[0])
-        XX, YY = np.meshgrid(xs, xs)
-        return np.array([XX, YY])
-    
-    def get_xycoords(self):
-        """Returns x_coords, y,_coords"""
-        return self.pixelscale * self.get_XXYY()
-    
-    
-    
-    
-    """Interpolation"""
-    def interp(self, coords, reim=False):
-        """ """
-        if not reim:
-            ampl  = map_coordinates(self.amplitude, coords, order=1)
-            phase = map_coordinates(self.phase,     coords, order=1)
-        else:
-            real = map_coordinates(self.get_real, coords, order=1)
-            imag = map_coordinates(self.get_imag, coords, order=1)
-            ampl = np.hypot(real, imag)
-            phase = np.arctan2(imag, real)
-        return ampl, phase
-        
-    def paraxial_interp(self, pixelscale_out, npix_out, reim=False):
-        """ Paraxially interpolates """
-        # Get coords arrays
-        npix_in = self.amplitude.shape[0]
-        ratio = pixelscale_out/self.pixelscale
-        
-        c_in = (npix_in-1)/2
-        c_out = (npix_out-1)/2
-        xs = ratio*(-c_out, c_out, npix_out) + c_in
-        YY, XX = np.meshgrid(xs, xs)
-        coords = np.array([XX, YY])
-        ampl, phase = self.interp(ampl, phase, reim=reim)
-        
-        # Update Phasor
-        self = self.update_phasor(ampl, phase)
-        
-        # Conserve energy
-        self = self.multiply_ampl(ratio)
-        
-        # Update pixelscale
-        return eqx.tree_at(lambda wf: wf.pixelscale, self, pixelscale_out)
-
-    
-    
-    """ Pad and Crop """
-    def pad_to(self, npix_out):
-        """ Pads an array paraxially to a given size
-        Works for even -> even or odd -> odd """
-        npix_in = self.amplitude.shape[0]
-        
-        if npix_in%2 != npix_out%2:
-            raise ValueError("Only supports even -> even or 0dd -> odd padding")
-        
-        c, s, rem = npix_out//2, npix_in//2, npix_in%2
-        padded = np.zeros([npix_out, npix_out])
-        
-        ampl  = padded.at[c-s:c+s+rem, c-s:c+s+rem].set(self.amplitude)
-        phase = padded.at[c-s:c+s+rem, c-s:c+s+rem].set(self.phase)
-        return self.update_phasor(ampl, phase)
 
-    def crop_to(self, npix_out):
-        """ Crops an array paraxially to a given size
-        Works for even -> even or odd -> odd """
-        npix_in = self.amplitude.shape[0]
-        
-        if npix_in%2 != npix_out%2:
-            raise ValueError("Only supports even -> even or 0dd -> odd cropping")
-        
-        c, s = npix_in//2, npix_out//2
-        ampl  = self.amplitude[c-s:c+s, c-s:c+s]
-        phase = self.phase[c-s:c+s, c-s:c+s]
-        return self.update_phasor(ampl, phase)
-        
-        
-        
-        
-    """Inversions"""
-    def invertXY(self):
-        ampl = self.amplitude[::-1, ::-1]
-        phase = self.phase[::-1, ::-1]
-        return self.update_phasor(ampl, phase)
-        
-    def invertX(self):
-        ampl = self.amplitude[:, ::-1]
-        phase = self.phase[:, ::-1]
-        return self.update_phasor(ampl, phase)
-
-    def invertY(self):
-        ampl = self.amplitude[::-1]
-        phase = self.phase[::-1]
-        return self.update_phasor(ampl, phase)
-    
-
-    
 class OpticalSystem(eqx.Module):
     """ Optical System class, Equinox Modle
     
diff --git a/src/propagators.py b/src/propagators.py
index 298441b2..6eaf953e 100644
--- a/src/propagators.py
+++ b/src/propagators.py
@@ -1,5 +1,11 @@
 import jax.numpy as np
 import equinox as eqx
+import typing
+
+
+GaussianPropagator = typing.NewType("GaussianPropgator", object)
+GaussianWavefront = typing.NewType("GaussianWavefront", object)
+
 
 class MFT(eqx.Module):
     """
@@ -231,4 +237,325 @@ def mft(self, wavefront, num_fringe, npix):
         t2 = np.dot(t1, expXU)
         wavefront_out = norm_coeff * t2
 
-        return wavefront_out
\ No newline at end of file
+        return wavefront_out
+
+
+class GaussianPropagator(eqx.Module):
+    """
+    An intermediate plane fresnel algorithm for propagating the
+    `GaussianWavefront` class between the planes. The propagator 
+    is separate from the `Wavefront` to emphasise the similarity 
+    of these algorithms to the layers in machine learning algorithms.
+
+    Attributes
+    ----------
+    distance : float 
+       The distance to propagate in meters. 
+    """
+    distance : float 
+
+
+    def __init__(self : GaussianPropagator, 
+            distance : float) -> GaussianPropagator:
+        """
+        Constructor for the GaussianPropagator.
+
+        Parameters
+        ----------
+        distance : float
+            The distance to propagate the wavefront in meters.
+        """
+        self.distance = distance
+
+
+    def planar_to_planar(self : GaussianPropagator, 
+            wavefront: GaussianWavefront, 
+            distance: float) -> GaussianWavefront:
+        """
+        Modifies the state of the wavefront by propagating a planar 
+        wavefront to a planar wavefront. 
+
+        Parameters
+        ----------
+        wavefront : GaussianWavefront
+            The wavefront to propagate. Must be `GaussianWavefront`
+            or a subclass. 
+        distance : float
+            The distance of the propagation in metres.
+
+        Returns
+        -------
+        : GaussianWavefront
+            The new `Wavefront` propagated by `distance`. 
+        """
+        complex_wavefront = wavefront.get_amplitude() * \
+            np.exp(1j * self.get_phase())
+
+        new_complex_wavefront = np.fft.ifft2(
+            wavefront.transfer_function(distance) * \
+            np.fft.fft2(complex_wavefront))
+
+        new_amplitude = np.abs(new_complex_wavefront)
+        new_phase = np.angle(new_complex_wavefront)
+        
+        return wavefront\
+            .set_position(wavefront.get_position() + distance)\
+            .update_phasor(new_amplitude, new_phase)        
+
+
+    def waist_to_spherical(self : GaussianPropagator, 
+            wavefront: GaussianWavefront, 
+            distance: float) -> GaussianWavefront:
+        """
+        Modifies the state of the wavefront by propagating it from 
+        the waist of the gaussian beam to a spherical wavefront. 
+
+        Parameters
+        ----------
+        wavefront : GaussianWavefront
+            The `Wavefront` that is getting propagated. Must be either 
+            a `GaussianWavefront` or a valid subclass.
+        distance : float 
+            The distance of the propagation in metres.
+
+        Returns 
+        -------
+        : GaussianWavefront
+            `wavefront` propgated by `distance`.
+        """
+        coefficient = 1 / 1j / wavefront.get_wavelength() / distance
+        complex_wavefront = wavefront.get_amplitude() * \
+            np.exp(1j * self.get_phase())
+
+        fourier_transform = jax.lax.cond(np.sign(distance) > 0, 
+            lambda wavefront, distance: \
+                quadratic_phase_factor(distance) * \
+                np.fft.fft2(wavefront), 
+            lambda wavefront, distance: \
+                quadratic_phase_factor(distance) * \
+                np.fft.ifft2(wavefront),
+            complex_wavefront, distance)
+
+        new_complex_wavefront = coefficient * fourier_transform
+        new_phase = np.angle(new_complex_wavefront)
+        new_amplitude = np.abs(new_complex_wavefront)
+
+        return wavefront\
+            .update_phasor(new_amplitude, new_phase)\
+            .set_position(wavefront.get_position() + distance)
+
+
+    def spherical_to_waist(self : GaussianPropagator, 
+            wavefront: GaussianWavefront, distance: float) -> GaussianWavefront:
+        """
+        Modifies the state of the wavefront by propagating it from 
+        a spherical wavefront to the waist of the Gaussian beam. 
+
+        Parameters
+        ----------
+        wavefront : GaussianWavefront 
+            The `Wavefront` that is getting propagated. Must be either 
+            a `GaussianWavefront` or a direct subclass.
+        distance : float
+            The distance of propagation in metres.
+
+        Returns
+        -------
+        : GaussianWavefront
+            The `wavefront` propagated by distance. 
+        """
+        # TODO: consider adding get_complex_wavefront() as this is 
+        # used in a number of locations. This could be implemented
+        # in the PhysicalWavefront and inherrited. 
+        coefficient = 1 / 1j / wavefront.get_wavelength() / \
+            distance * wavefront.quadratic_phase_factor(distance)
+        complex_wavefront = wavefront.get_amplitude() * \
+            np.exp(1j * wavefront.get_phase())
+
+        fourier_transform = jax.lax.cond(np.sign(distance) > 0, 
+            lambda wavefront: np.fft.fft2(wavefront), 
+            lambda wavefront: np.fft.ifft2(wavefront),
+            complex_wavefront)
+
+        new_wavefront = coefficient * fourier_transform
+        new_phase = np.angle(new_wavefront)
+        new_amplitude = np.abs(new_wavefront)
+        return wavefront\
+            .update_phasor(new_amplitude, new_phase)\
+            .set_position(wavefront.get_position() + distance)
+
+
+    def outside_to_outside(self : GaussianWavefront, 
+            wavefront : GaussianWavefront, 
+            distance: float) -> GaussianWavefront:
+        """
+        Propagation from outside the Rayleigh range to another 
+        position outside the Rayleigh range. 
+
+        Parameters
+        ----------
+        wavefront : GaussianWavefront
+            The wavefront to propagate. Assumed to be either a 
+            `GaussianWavefront` or a direct subclass.
+        distance : float
+            The distance to propagate in metres.
+    
+        Returns
+        -------
+        : GaussianWavefront 
+            The new `Wavefront` propgated by distance. 
+        """
+        # TODO: Make sure that the displacements are called by the 
+        # correct function.
+        from_waist_displacement = wavefront.get_position() \
+            + distance - wavefront.location_of_waist()
+        to_waist_displacement = wavefront.location_of_waist() \
+            - wavefront.get_position()
+
+        wavefront_at_waist = self.spherical_to_waist(
+            wavefront, to_waist_displacement)
+        wavefront_at_distance = self.spherical_to_waist(
+            wavefront_at_waist, from_waist_displacement)
+
+        return wavefront_at_distance
+
+
+    def outside_to_inside(self : GaussianPropagator,
+            wavefront: GaussianWavefront, distance: float) -> None:
+        """
+        Propagation from outside the Rayleigh range to inside the 
+        rayleigh range. 
+
+        Parameters
+        ----------
+        wavefront : GaussianWavefront
+            The wavefront to propagate. Must be either 
+            `GaussianWavefront` or a direct subclass.
+        distance : float
+            The distance to propagate in metres.
+
+        Returns
+        -------
+        : GaussianWavefront
+            The `Wavefront` propagated by `distance` 
+        """
+        from_waist_displacement = wavefront.position + distance - \
+            wavefront.location_of_waist()
+        to_waist_position = wavefront.location_of_waist() - \
+            wavefront.get_position()
+
+        wavefront_at_waist = self.planar_to_planar(
+            wavefront, to_waist_displacement)
+        wavefront_at_distance = self.spherical_to_waist(
+            wavefront_at_waist, from_waist_displacement)
+
+        return wavefront_at_distance
+
+
+    def inside_to_inside(self : GaussianPropagator, 
+            wavefront : GaussianWavefront, 
+            distance: float) -> GaussianWavefront:
+        """
+        Propagation from inside the Rayleigh range to inside the 
+        rayleigh range. 
+
+        Parameters
+        ----------
+        wavefront : GaussianWavefront
+            The `Wavefront` to propagate. This must be either a 
+            `GaussianWavefront` or a direct subclass.
+        distance : float
+            The distance to propagate in metres.
+
+
+        Returns
+        -------
+        : GaussianWavefront
+            The `Wavefront` propagated by `distance`
+        """
+        # TODO: Consider removing after checking the Jaxpr for speed. 
+        # This is just an alias for another function. 
+        return self.planar_to_planar(wavefront, distance)
+
+
+    def inside_to_outside(self : GaussianPropagator, 
+            wavefront : GaussianWavefront, 
+            distance: float) -> GaussianWavefront:
+        """
+        Propagation from inside the Rayleigh range to outside the 
+        rayleigh range. 
+
+        Parameters
+        ----------
+        wavefront : GaussianWavefront
+            The `Wavefront` to propgate. Must be either a 
+            `GaussianWavefront` or a direct subclass.
+        distance : float
+            The distance to propagate in metres.
+
+        Returns
+        -------
+        : GaussianWavefront 
+            The `Wavefront` propagated `distance`.
+        """
+        from_waist_displacement = wavefront.get_position() + \
+            distance - wavefront.location_of_waist()
+        to_waist_displacement = wavefront.location_of_waist() - \
+            wavefront.get_position()
+
+        wavefront_at_waist = self.planar_to_planar(
+            wavefront, to_waist_displacement)
+        wavefront_at_distance = self.waist_to_spherical(
+            wavefront_at_waist, from_waist_displacement)
+
+        return wavefront_at_distance
+
+
+    def __call__(self : GaussianPropagator, 
+            wavefront : GaussianWavefront, 
+            distance: float) -> GaussianWavefront:
+        """
+        Propagates the wavefront approximated by a Gaussian beam 
+        the amount specified by distance. Note that distance can 
+        be negative.
+
+        Parameters 
+        ----------
+        wavefront : GaussianWavefront
+            The `Wavefront` to propagate. Must be a `GaussianWavefront`
+            or a direct subclass.
+        distance : float
+            The distance to propagate in metres.
+
+        Returns
+        -------
+        : GaussianWavefront 
+            The `Wavefront` propagated `distance`.
+        """
+        # This works by considering the current position and distnace 
+        # as a boolean array. The INDEX_GENERATOR converts this to 
+        # and index according to the following table.
+        #
+        # sum((0, 0) * (1, 2)) == 0
+        # sum((1, 0) * (1, 2)) == 1
+        # sum((0, 1) * (1, 2)) == 2
+        # sum((1, 1) * (1, 2)) == 3
+        #
+        # TODO: Test if a simple lookup is faster. 
+        # Constants
+        INDEX_GENERATOR = np.array([1, 2])
+
+        decision_vector = wavefront.is_inside([
+            wavefront.get_position(), 
+            wavefront.get_position() + distance])
+        decision_index = np.sum(
+            self.INDEX_GENERATOR * descision_vector)
+ 
+        # Enters the correct function differentiably depending on 
+        # the descision.
+        new_wavefront = jax.lax.switch(decision_index, 
+            [self.inside_to_inside, self.inside_to_outside,
+            self.outside_to_inside, self.outside_to_outside],
+            wavefront, distance) 
+
+        return new_wavefront
diff --git a/src/wavefronts.py b/src/wavefronts.py
new file mode 100644
index 00000000..d3c9d161
--- /dev/null
+++ b/src/wavefronts.py
@@ -0,0 +1,1121 @@
+import equinox as eqx 
+import jax.numpy as np
+import typing
+
+
+Wavefront = typing.NewType("Wavefront", eqx.Module)
+PhysicalWavefront = typing.NewType("PhysicalWavefront", Wavefront)
+AngularWavefront = typing.NewType("AngularWavefront", Wavefront)
+GaussianWavefront = typing.NewType("FresnelWavefront", Wavefront)
+Array = typing.NewType("Array", np.ndarray)
+
+
+class Wavefront(eqx.Module):
+    """
+    An abstract module that should never be directly substantiated.
+    This class represents a general optical wavefront although the 
+    exact implementation must go through a subclass. Wavefront
+    objects are assumed to be square.
+
+
+    The Wavefront and its are not intended to be public functionality
+    and are initialised in a vectorised manner by the CreateWavefront 
+    layer which represents the source of the optical disturbance. 
+
+
+    Attributes
+    ----------
+    wavelength : float
+        The wavelength of the light. Assumed to be in metres.
+    offset : Array
+        The polarisation state of the wave described by the x and 
+        y phase lag of the wavefront. This quantity is unitless and 
+        it is assumed that `offset.shape == (2, )`  
+    """
+    wavelength : float
+    offset : Array
+    plane_type : str # For debugging
+    amplitude : Array
+    phase : Array
+    pixel_scale : float
+
+
+    def __init__(self : Wavefront, wavelength : float, 
+            offset : Array) -> Wavefront:
+        """
+        Initialises a minimal wavefront specified only by the 
+        wavel. 
+
+        Parameters
+        ----------
+        wavelength : float 
+            The monochromatic wavelength associated with this 
+            wavefront. 
+        offset : Array
+            The x and y angles of incidence to the surface assumed to 
+            be in radians. 
+        """
+        self.wavelength = np.array([wavelength]).astype(float) # Jax Safe
+        self.offset = np.array([offset]).astype(float) # To be instantiated by CreateWavefront        
+        self.plane_type = "Pupil"
+        self.amplitude = None
+        self.phase = None
+        self.pixel_scale = None
+
+
+    def get_pixel_scale(self : GaussianWavefront) -> float:
+        """
+        Accessor for the pixel_scale.
+
+        Returns
+        -------
+        : float
+            The pixel_scale associated with the current position.
+        """
+        return self.pixel_scale
+
+
+    def get_offset(self : Wavefront) -> Array:
+        """
+        The offset of this `Wavefront`.
+
+        Returns
+        -------
+        : Array
+            The x and y angles of deviation that the wavefront makes
+            from the mirror.
+        """
+        return self.offset
+
+
+    def set_offset(self : Wavefront, offset : Array) -> Wavefront:
+        """
+        Mutator for the offset.
+
+        Parameters
+        ----------
+        offset : Array (f64[2])
+            The angles that the `Wavefront` makes with the x and y 
+            axis assumed to be in radians.
+        
+        Returns
+        -------
+        : Wavefront
+            The new `Wavefront` with the updated offset. 
+        """
+        return eqx.tree_at(
+            lambda wavefront : wavefront.offset, self, offset,
+            is_leaf = lambda leaf : leaf is None )
+
+
+    def get_wavelength(self : Wavefront) -> float:
+        """
+        Accessor for the wavelength.
+
+        Returns
+        -------
+        : float
+            The wavelength of the `Wavefront` in meters. 
+        """
+        return self.wavelength
+
+
+    def set_wavelength(self : Wavefront, wavelength : float) -> Wavefront:
+        """
+        Mutator for the wavelength.
+
+        Parameters
+        ----------
+        wavelength : float  
+            The new wavelength of the `Wavefront` assumed to be in meters.
+
+        Returns
+        -------
+        : Wavefront
+            The new `Wavefront` with the updated wavelength. 
+        """
+        return eqx.tree_at(
+            lambda wavefront : wavefront.wavelength, self, wavelength,
+            is_leaf = lambda leaf : leaf is None)
+
+
+    def get_amplitude(self : Wavefront) -> Array:
+        """
+        Accessor for the amplitude.
+
+        Returns
+        -------
+        : Array 
+            The electric field amplitude in SI units of electric field. 
+        """
+        return self.amplitude
+
+
+    def get_phase(self : Wavefront) -> Array:
+        """
+        Accessor for the phase.
+
+        Returns
+        -------
+        : Array 
+            The phase of the Wavefront; a unitless qunatity.
+        """
+        return self.phase
+
+
+    def set_amplitude(self : Wavefront, amplitude : Array) -> Wavefront:
+        """
+        Mutator for the amplitude. 
+
+        Parameters
+        ---------
+        amplitude : Array
+            The new amplitudes of the `Wavefront` assumed to be in the 
+            SI units of electric field.         
+
+        Returns
+        -------
+        : Wavefront
+            The new `Wavefront` with the updated amplitude. 
+        """
+        return eqx.tree_at(
+            lambda wavefront : wavefront.amplitude, self, amplitude,
+            is_leaf = lambda leaf : leaf is None)
+
+
+    def set_phase(self : Wavefront, phase : Array) -> Wavefront:
+        """
+        Mutator for the phase.
+
+        Parameters
+        ----------
+        phase : Array
+            The new phases of the `Wavefront` assumed to be unitless.
+
+        Returns
+        -------
+        : Wavefront
+            The new `Wavefront` with `Wavefront.get_phase() == phase`.
+        """
+        return eqx.tree_at(
+            lambda wavefront : wavefront.phase, self, phase,
+            is_leaf = lambda leaf : leaf is None) 
+
+
+    def get_real(self : Wavefront) -> Array:
+        """
+        The real component of the `Wavefront`. 
+
+        Throws
+        ------
+        : TypeError
+            If self.amplitude or self.phase have not been initialised
+            externally.  
+
+        Returns 
+        -------
+        : Array
+            The real component of the optical disturbance with 
+            SI units of electric field.  
+        """
+        return self.get_amplitude() * np.cos(self.get_phase())
+        
+
+    def get_imaginary(self : Wavefront) -> Array:
+        """
+        The imaginary component of the `Wavefront`
+
+        Throws
+        ------
+        : TypeError
+            If self.amplitude or self.phase have not been initialised
+            externally. 
+
+        Returns
+        -------
+        : Array
+            The imaginary component of the optical disturbance with 
+            the SI units of electric field. 
+        """
+        return self.get_amplitude() * np.sin(self.get_phase())
+
+
+    def multiply_amplitude(self : Wavefront, 
+            array_like : typing.Union[float, Array]) -> Wavefront:
+        """
+        Modify the amplitude of the wavefront via elementwise 
+        multiplication. 
+
+        Throws
+        ------
+        : TypeError
+            If `self.amplitude` has not initialised externally.
+        : ValueError
+            If `weights` is not a scalar, or an array of the same 
+            dimensions. i.e.
+            ```py
+            ((weights.shape == (1,)) \
+                | (weights.shape == self.amplitude.shape) == True
+            ```
+
+        Parameters
+        ----------
+        weights : Union[float, array]
+            An array that has the same dimensions as self.amplitude 
+            by which elementwise multiply each pixel. 
+            A float to apply to the entire array at once. May simulate 
+            transmission through a translucent element.
+
+        Returns
+        -------
+        : Wavefront
+            The new Wavefront with the applied changes to the 
+            amplitude array. 
+        """
+        return self.set_amplitude(self.get_amplitude() * array_like)
+
+
+    def add_phase(self : Wavefront, 
+            phases : typing.Union[float, Array]) -> Wavefront:
+        """
+        Used to update the wavefront phases based on the current 
+        position using elementwise addition. 
+
+        Throws
+        ------
+        : TypeError
+            If self.phase has not been initialised externally.
+        : ValueError
+            If `amounts` is not of the same dimensions as `self.phases`
+            or `amounts`. i.e. 
+            ```py
+            ((weights.shape == (1,)) \
+                | (weights.shape == self.amplitude.shape) == True
+            ```
+
+        Parameters
+        ----------
+        amounts : Union[float, array]
+            The amount of phase to add to the current phase value of 
+            each pixel. A scalar modifies the global phase of the 
+            wavefront. 
+
+        Returns
+        -------
+        : Wavefront
+            The new wavefront with the updated array of phases. 
+        """
+        return self.set_phase(self.get_phase() + phases)
+
+
+    def update_phasor(self : Wavefront, amplitude : Array, 
+            phase : Array) -> Wavefront:  
+        """
+        Used to write the state of the optical disturbance. Ignores 
+        the current state. It is assumed that `amplitude` and `phase`
+        have the same shape i.e. `amplitude.shape == phase.shape`.
+        It is not assumed that the shape of the wavefront is 
+        maintained i.e. `self.amplitude.shape == amplitude.shape`
+        is __not__ required. 
+
+        Parameters
+        ----------
+        amplitude : Array
+            The electric field amplitudes of the wavefront. Assumed to
+            have the SI units of electric field. 
+        phase : Array
+            The phases of each pixel in the new wavefront. Assumed to 
+            be unitless.
+
+        Returns
+        -------
+        : Wavefront
+            The new wavefront with specified by `amplitude` and `phase`        
+        """
+        return self.set_phase(phase).set_amplitude(amplitude)
+
+
+    def wavefront_to_psf(self : Wavefront) -> Array:
+        """
+        Calculates the _P_oint _S_pread _F_unction (PSF) of the 
+        wavefront. 
+
+        Throws
+        ------
+        : TypeError
+            If `self.amplitude` has not been externally initialised.
+
+        Returns
+        -------
+        : Array
+            The PSF of the wavefront.
+        """
+        return self.get_amplitude() ** 2
+
+
+    def add_opd(self: Wavefront, 
+            path_difference : typing.Union[float, Array]) -> Wavefront:
+        """
+        Changes the state of the wavefront based on the optical path 
+        taken. 
+
+        Throws
+        ------
+        : TypeError
+            If `self.phase` has not been externally initialised
+        : ValueError
+            If `path_difference.shape != self.phase.shape` or 
+            If `path_difference.shape != (1,)`
+
+        Parameters
+        ----------
+        path_difference : Union[float, Array]
+            The physical path difference in meters of either the 
+            entire wavefront or each pixel individually. 
+        
+        Returns
+        -------
+        : Wavefront
+            The new wavefront with the phases updated according to 
+            `path_difference`     
+        """
+        phase_difference = 2 * np.pi * path_difference / self.wavelength
+        return self.add_phase(phase_difference)
+
+
+    def normalise(self : Wavefront) -> Wavefront:
+        """
+        Reduce the electric field amplitude of the wavefront to a 
+        range between 0 and 1. Guarantees that:
+        
+        Throws
+        ------
+        : TypeError
+            If `self.amplitude` has not been externally instantiated.
+
+        Returns
+        -------
+        : Wavefront
+            The new wavefront with the normalised electric field 
+            amplitudes. The amplitude is now unitless. 
+        """
+        total_intensity = np.linalg.norm(self.get_amplitude())
+        return self.multiply_amplitude(1 / total_intensity)
+
+
+    def get_pixel_coordinates(self : Wavefront, 
+            number_of_pixels : int) -> Array:
+        """
+        A static helper method for correctly creating paraxially pixel 
+        arrays for optical transformations. 
+
+        Parameters
+        ----------
+        number_of_pixels : int 
+            The length of the paraxial pixel array.
+
+        Returns
+        -------
+        : Array
+            The paraxial pixel positions of with dimensions 
+            `number_of_pixels`
+        """
+        return np.arange(number_of_pixels) - (number_of_pixels - 1) / 2
+
+
+    def get_pixel_grid(self : Wavefront) -> Array:
+        """
+        The pixel positions corresponding to each entry in the 
+        optical disturbances stored in the Wavefront. 
+
+        Throws
+        ------
+        : TypeError
+            If `self.amplitude` has not been externally initialised
+
+        Returns
+        -------
+        : Array 
+            The pixel positions of the optical disturbances. 
+            Guarantees `self.get_pixel_grid().shape == 
+            self.amplitude.shape`
+        """
+        pixel_positions = self.get_pixel_coordinates(self.amplitude.shape[0])
+        x_positions, y_positions = \
+            np.meshgrid(pixel_positions, pixel_positions)
+        return np.array([x_positions, y_positions])
+
+
+    def get_pixel_positions(self : Wavefront) -> Array:
+        """
+        The physical coordinates of each optical disturbance in meters.
+
+        Throws
+        ------
+        : TypeError
+            If `self.amplitude` has not been externally initialised
+        : TypeError
+            If `self.pixelscale` has not been externally initialised
+
+        Returns
+        -------
+        : Array
+            The physical positions of the optical disturbance. 
+            Guarantees that `self.get_coordinates().shape == 
+            self.amplitude.shape`.
+        """
+        return self.get_pixel_scale() * self.get_pixel_grid()
+
+
+    def invert_x_and_y(self : Wavefront) -> Wavefront:
+        """
+        Reflects the wavefront across both axes. 
+
+        Throws
+        ------
+        : ValueError
+            If `self.amplitude` is not externally initialised.
+        : ValueError
+            If `self.phase` not externally initialised.
+
+        Returns
+        -------
+        : Wavefront
+            The new `Wavefront` with the phase and amplitude arrays
+            reversed accros both axes.
+        """
+        return self.invert_x().invert_y()
+
+
+    def invert_x(self : Wavefront) -> Wavefront:
+        """
+        Reflects the wavefront across the x axis.
+
+        Throws
+        ------
+        : ValueError
+            If `self.amplitude` is not externally initialised.
+        : ValueError 
+            If `self.phase` is not externally initialised. 
+
+        Returns
+        -------
+        : Wavefront
+            The new `Wavefront` with the phase and the amplitude arrays
+            reversed along the x axis. 
+        """
+        new_amplitude = self.amplitude[:, ::-1]
+        new_phase = self.phase[:, ::-1]
+        return self.update_phasor(new_amplitude, new_phase)
+
+
+    def invert_y(self : Wavefront) -> Wavefront:
+        """
+        Reflects the wavefront across the y axis.
+
+        Throws
+        ------        
+        : ValueError
+            If `self.amplitude` is not externally initialised.
+        : ValueError 
+            If `self.phase` is not externally initialised. 
+
+        Returns
+        -------
+        : Wavefront
+            The new wavefront with the phase and the amplitude arrays 
+            reversed along the y axis.
+        """
+        new_amplitude = self.amplitude[::-1, :]
+        new_phase = self.phase[::-1, :]
+        return self.update_phasor(new_amplitude, new_phase)
+
+
+    # TODO: Make logic Jax-Safe
+    def interpolate(self : PhysicalWavefront, coordinates : Array, 
+            real_imaginary : bool = False) -> tuple[Array, Array]:
+        """
+        Interpolates the `Wavefront` at the points specified by 
+        coordinates. The default interpolation uses the amplitude 
+        and phase although by passing `real_imgainary == True` 
+        the interpolation can be based on the real and imaginary 
+        information. The main use of this function is as a helper 
+        method to `self.paraxial_interpolate`.
+
+        Parameters
+        ----------
+        coordinates : Array
+            The coordinates to interpolate the optical disturbance
+            at. Assumed to have the units meters. 
+        real_imaginary : bool
+            Whether to use the amplitude-phase or real-imaginary
+            representation for the interpolation. The amplitude-
+            phase representation is slightly faster.
+
+        Returns
+        -------
+        : tuple[Array, Array]
+            The amplitude and phase of the wavefront at `coordinates`
+            based on a linear interpolation.
+        """
+        if not real_imaginary:
+            new_amplitude = map_coordinates(
+                self.amplitude, coordinates, order=1)
+            new_phase = map_coordinates(
+                self.phase, coordinates, order=1)
+        else:
+            real = map_coordinates(
+                self.get_real(), coordinates, order=1)
+            imaginary = map_coordinates(
+                self.get_imaginary(), coordinates, order=1)
+            new_amplitude = np.hypot(real, imaginary)
+            new_phase = np.arctan2(imaginary, real)
+        return new_amplitude, new_phase
+
+
+    def paraxial_interpolate(self : PhysicalWavefront, 
+            pixel_scale_out : float, number_of_pixels_out : int,
+            real_imaginary : bool = False) -> PhysicalWavefront: 
+        """
+        Interpolates the `Wavefront` so that it remains centered on 
+        each pixel (paraxial). Calculation can be performed using 
+        either the real-imaginary or amplitude-phase representations 
+        of the wavefront. The default is amplitude-phase. 
+
+        Parameters
+        ----------
+        pixel_scale_out : float
+            The dimensions of a single square pixel after the 
+            interpolation.
+        number_of_pixels_out : int
+            The number of pixels along one side of the square
+            `Wavefront` after the interpolation. 
+        real_imaginary : bool = False
+            Whether to use the real-imaginary representation of the 
+            wavefront for the interpolation. 
+
+        Returns
+        -------
+        : PhysicalWavefront
+            The new wavefront with the updated optical disturbance. 
+        """
+        # Get coords arrays
+        number_of_pixels_in = self.amplitude.shape[0]
+        ratio = pixel_scale_out / self.pixel_scale
+        
+        centre = (number_of_pixels_in - 1) / 2
+        new_centre = (number_of_pixels_out - 1) / 2
+        pixels = ratio * (-new_centre, new_centre, number_of_pixels_out) + centre
+        x_pixels, y_pixels = np.meshgrid(pixels, pixels)
+        coordinates = np.array([y_pixels, x_pixels])
+        new_amplitude, new_phase = self.interpolate(
+            coordinates, real_imaginary=real_imaginary)
+        
+        # Update Phasor
+        self = self.update_phasor(new_amplitude, new_phase)
+        
+        # Conserve energy
+        self = self.multiply_amplitude(ratio)
+        
+        return self.set_pixel_scale(pixel_scale_out)
+
+
+    def set_pixel_scale(self : Wavefront, pixel_scale : float) -> Wavefront:
+        """
+        Mutator for the pixel_scale.
+
+        Parameters
+        ----------
+        pixel_scale : float
+            The new pixel_scale associated with the wavefront.
+
+        Returns
+        -------
+        : Wavefront
+            The new Wavefront object with the updated pixel_scale
+        """
+        return eqx.tree_at(
+            lambda wavefront : wavefront.pixel_scale, self, pixel_scale,
+            is_leaf = lambda leaf : leaf is None)
+
+
+    def pad_to(self : PhysicalWavefront, 
+            number_of_pixels_out : int) -> PhysicalWavefront:
+        """
+        Pads the `Wavefront` with zeros. Assumes that 
+        `number_of_pixels_out > self.amplitude.shape[0]`. 
+        Note that `Wavefronts` with even pixel dimensions can 
+        only be padded (without interpolation) to even pixel 
+        dimensions and vice-versa. 
+
+        Throws
+        ------
+        : ValueError
+            If `number_of_pixels_out%2 != self.amplitude.shape[0]%2`
+            i.e. padding an even (odd) `Wavefront` to odd (even).
+
+        Parameters
+        ----------
+        number_of_pixels_out : int
+            The square side length of the array after it has been 
+            zero padded. 
+
+
+        Returns
+        -------
+        : PhysicalWavefront
+            The new `Wavefront` with the optical disturbance zero 
+            padded to the new dimensions.
+        """
+        number_of_pixels_in = self.amplitude.shape[0]
+
+        if number_of_pixels_in % 2 != number_of_pixels_out % 2:
+            raise ValueError("Only supports even -> even or odd -> odd padding")
+       
+        # TODO: Error (>) for the JAX silent error. 
+ 
+        new_centre = number_of_pixels_out // 2
+        centre = number_of_pixels_in // 2
+        remainder = number_of_pixels_in % 2
+        padded = np.zeros([number_of_pixels_out, number_of_pixels_out])
+        
+        new_amplitude = padded.at[
+                new_centre - centre : centre + new_centre + remainder, 
+                new_centre - centre : centre + new_centre + remainder
+            ].set(self.amplitude)
+        new_phase = padded.at[
+                new_centre - centre : centre + new_centre + remainder, 
+                new_centre - centre : centre + new_centre + remainder
+            ].set(self.phase)
+        return self.update_phasor(new_amplitude, new_phase)
+
+
+    def crop_to(self : PhysicalWavefront, 
+            number_of_pixels_out : int) -> PhysicalWavefront:
+        """
+        Crops a `Wavefront`'s optical disturbance. Assumes that 
+        `number_of_pixels_out < self.amplitude.shape[0]`. 
+        `Wavefront`s with an even number of pixels can only 
+        be cropped to an even number of pixels without interpolation
+        and vice-versa.    
+        
+        Throws
+        ------
+        : ValueError
+            If `number_of_pixels_out%2 != self.amplitude.shape[0]%2`
+            i.e. padding an even (odd) `Wavefront` to odd (even).
+
+        Parameters
+        ----------
+        number_of_pixels_out : int
+            The square side length of the array after it has been 
+            zero padded. 
+
+
+        Returns
+        -------
+        : PhysicalWavefront
+            The new `Wavefront` with the optical disturbance zero 
+            cropped to the new dimensions.
+        """
+        number_of_pixels_in = self.amplitude.shape[0]
+        
+        if number_of_pixels_in%2 != number_of_pixels_out%2:
+            raise ValueError("Only supports even -> even or 0dd -> odd cropping")
+        
+        new_centre = number_of_pixels_in // 2
+        centre = number_of_pixels_out // 2
+
+        new_amplitude = self.amplitude[
+            new_centre - centre : new_centre + centre, 
+            new_centre - centre : new_centre + centre]
+        new_phase = self.phase[
+            new_centre - centre : new_centre + centre, 
+            new_centre - centre : new_centre + centre]
+
+        return self.update_phasor(new_amplitude, new_phase)
+
+
+class PhysicalWavefront(Wavefront):
+    """
+    A simple plane wave extending the abstract `Wavefront` class. 
+    Assumes that the wavefront is square. This is Physical as 
+    opposed to Angular, because there are no infinite distances.  
+
+    Attributes
+    ----------
+    plane_type : str
+        The type of plane occupied by the wavefront. 
+    amplitude : Array
+        The electric field amplitudes over the wavefront. The 
+        amplitude is assumed to be in SI units. 
+    phase : Array
+        The phases of each pixel on the Wavefront. The phases are 
+        assumed to be unitless.
+    pixel_scale : float
+        The physical dimensions of each square pixel. Assumed to be 
+        metres.
+    offset : Array
+        The angle that the `Wavefront` makes with the OpticalElement
+        in radians. This is a (1, 2) array such that 
+        `offset <= 2 * np.pi`. 
+    """
+    def __init__(self : PhysicalWavefront, wavelength : float, 
+            offset : Array) -> PhysicalWavefront:
+        """
+        Constructor for a `PhysicalWavefront`.
+
+        Parameters
+        ----------
+        wavelength : float 
+            The wavelength associated with this wavefront in 
+            meters.
+        offset : Array
+            The angle that the wavefront makes with the element
+            in the x and y coordinates.
+
+        Returns
+        -------
+        : PhysicalWavefront
+            The new wavefront with `None` at the extra leaves. 
+        """
+        super().__init__(wavelength, offset)
+
+
+class AngularWavefront(Wavefront):
+    """
+    A wavefront parametrised by phase and amplitude arrays, as 
+    well as the pixel scale. The units are radians as opposed to 
+    meters (in `PhysicalWavefront`).
+
+    Attributes
+    ----------
+    plane_type : str
+        The type of plane occupied by the wavefront. 
+    amplitude : Array
+        The electric field amplitudes over the wavefront. The 
+        amplitude is assumed to be in SI units. 
+    phase : Array
+        The phases of each pixel on the Wavefront. The phases are 
+        assumed to be unitless.
+    pixel_scale : float
+        The physical dimensions of each square pixel. Assumed to be 
+        metres.
+    """ 
+    def __init__(self : AngularWavefront, wavelength : float, 
+            offset : Array) -> AngularWavefront:
+        """
+        Constructor for a `AngularWavefront`.
+
+        Parameters
+        ----------
+        wavelength : float 
+            The wavelength associated with this wavefront in 
+            meters.
+        offset : Array
+            The angle that the wavefront makes with the element
+            in the x and y coordinates.
+
+        Returns
+        -------
+        : AngularWavefront
+            The new wavefront with `None` at the extra leaves. 
+        """
+        super().__init__(wavelength, offset)
+
+
+class GaussianWavefront(Wavefront):
+    """
+    Expresses the behaviour and state of a wavefront propagating in 
+    an optical system under the fresnel assumptions. This 
+    implementation is based on the same class from the `poppy` 
+    library [poppy](https://github.com/spacetelescope/poppy/fresnel.py)
+    and Chapter 3 from _Applied Optics and Optical Engineering_
+    by Lawrence G. N.
+
+
+    Approximates the wavefront as a Gaussian Beam parametrised by the 
+    radius of the beam, the phase radius, the phase factor and the 
+    Rayleigh distance. Propagation is based on two different regimes 
+    for a total of four different opertations. 
+    
+    Attributes
+    ----------
+    position : float
+        The position of the wavefront in the optical system assumed 
+        to be in meters.
+    beam_radius : float
+        The radius of the beam assumed to be in metres. 
+    phase_radius : float
+        The phase radius of the gaussian beam assumed to be unitless. 
+    """
+    position : float 
+    beam_radius : float
+    phase_radius : float
+
+
+    def __init__(self : GaussianWavefront, 
+            offset : Array,
+            wavelength : float) -> GaussianWavefront:
+        """
+        Creates a wavefront with an empty amplitude and phase 
+        arrays but of a given wavelength and phase offset. 
+
+        Parameters
+        ----------
+        beam_radius : float
+            Radius of the beam at the initial optical plane.
+        wavelength : float
+            Wavelength of the monochromatic light.
+        offset : Array
+            Phase shift of the initial optical plane. 
+        phase_radius :  float
+            The phase radius of the GuasianWavefront. This is a unitless
+            quantity. 
+        """
+        super().__init__(wavelength, offset)
+        self.beam_radius = None
+        self.phase_radius = np.inf
+        self.position = None 
+
+
+    def get_position(self : GaussianWavefront) -> float:
+        """
+        Accessor for the position of the wavefront. 
+
+        Returns 
+        -------
+        : float 
+            The position of the `Wavefront` from its starting point 
+            in meters.
+        """
+        return self.position
+
+
+    def set_position(self : GaussianWavefront, 
+            position : float) -> GaussianWavefront:
+        """
+        Mutator for the position of the wavefront. Changes the 
+        pixel_scale which is a function of the position.  
+
+        Parameters
+        ----------
+        position : float
+            The new position of the wavefront from its starting point 
+            assumed to be in meters. 
+        
+        Returns
+        -------
+        : GaussianWavefront
+            This wavefront at the new position. 
+        """
+        return eqx.tree_at(
+            lambda wavefront : wavefront.position, self, position,
+            is_leaf = lambda leaf : leaf is None)
+
+
+    def get_beam_radius(self : GaussianWavefront) -> float:
+        """
+        Accessor for the radius of the wavefront.
+
+        Returns
+        -------
+        : float
+            The radius of the `GaussianWavefront` in meters.
+        """
+        return self.beam_radius
+
+
+    def get_phase_radius(self : GaussianWavefront) -> float:
+        """
+        Accessor for the phase radius of the wavefront.
+
+        Returns
+        -------
+        : float 
+            The phase radius of the wavefront. This is a unitless 
+            quantity.
+        """
+        return self.phase_radius
+
+
+    def rayleigh_distance(self: GaussianWavefront) -> float:
+        """
+        Calculates the rayleigh distance of the Gaussian beam.
+        
+        Returns
+        -------
+        : float
+            The Rayleigh distance of the wavefront in metres.
+        """
+        return np.pi * self.get_beam_radius() ** 2\
+            / self.get_wavelength()
+
+
+    def transfer_function(self: GaussianWavefront, distance: float) -> Array:
+        """
+        The optical transfer function (OTF) for the gaussian beam.
+        Assumes propagation is along the axis. 
+
+        Parameters
+        ----------
+        distance : float
+            The distance to propagate the wavefront along the beam 
+            via the optical transfer function in metres.
+
+        Returns
+        -------
+        : float 
+            A phase representing the evolution of the wavefront over 
+            the distance. 
+
+        References
+        ----------
+        Wikipedia contributors. (2022, January 3). Direction cosine. 
+        In Wikipedia, The Free Encyclopedia. June 23, 2022, from 
+        https://en.wikipedia.org/wiki/Direction_cosine
+
+        Wikipedia contributors. (2022, January 3). Spatial frequecy.
+        In Wikipedia, The Free Encyclopedia. June 23, 2022, from 
+        https://en.wikipedia.org/wiki/Spatial_frequency
+        """
+        coordinates = self.get_pixel_positions()
+        radius = np.sqrt((coordinates ** 2).sum(axis=0))
+        xi = coordinates[0, :, :] / radius / self.get_wavelength()
+        eta = coordinates[1, :, :] / radius / self.get_wavelength()
+        return np.exp(1j * np.pi * self.get_wavelength() \
+            * distance * (xi ** 2 + eta ** 2))
+
+
+    def quadratic_phase_factor(self: GaussianWavefront, 
+            distance: float) -> float:
+        """
+        Convinience function that simplifies many of the diffraction
+        equations. Caclulates a quadratic phase factor associated with 
+        the beam. 
+
+        Parameters
+        ----------
+        distance : float
+            The distance of the propagation measured in metres. 
+
+        Returns
+        -------
+        : float
+            The near-field quadratic phase accumulated by the beam
+            from a propagation of distance.
+        """      
+        return np.exp(1j * np.pi * \
+            (self.get_pixel_positions() ** 2).sum(axis=0) \
+            / self.get_wavelength() / distance)
+
+
+    def location_of_waist(self: GaussianWavefront) -> float:
+        """
+        Calculates the position of the waist along the direction of 
+        propagation based of the current state of the wave.
+
+        Returns
+        -------
+        : float
+            The position of the waist in metres.
+        """
+        return - self.get_phase_radius() / \
+            (1 + (self.get_phase_radius() / \
+            self.rayleigh_distance()) ** 2)
+
+
+    def waist_radius(self: GaussianWavefront) -> float:
+        """
+        The radius of the beam at the waist.
+
+        Returns
+        -------
+        : float
+            The radius of the beam at the waist in metres.
+        """
+        return self.get_beam_radius() / \
+            np.sqrt(1 + (self.rayleigh_distance() \
+                / self.get_beam_radius()) ** 2) 
+
+
+    def calculate_pixel_scale(self: GaussianWavefront, position: float) -> None:
+        """
+        The pixel scale at the position along the axis of propagation.
+        Assumes that the wavefront is square. That is:
+        ```
+        x, y = self.amplitude.shape
+        (x == y) == True
+        ```
+
+        Parameters
+        ----------
+        position : float
+            The position of the wavefront aling the axis of propagation
+            in metres.
+        """
+        # TODO: get_number_of_pixels() Used frequenctly not a function
+        number_of_pixels = self.amplitude.shape[0]
+        new_pixel_scale = self.get_wavelength() * np.abs(position) / \
+            number_of_pixels / self.get_pixel_scale()  
+        return new_pixel_scale 
+        
+
+    def is_inside(self: GaussianWavefront, distance: float) -> bool:
+        """ 
+        Determines whether a point at along the axis of propagation 
+        at distance away from the current position is inside the 
+        rayleigh distance. 
+
+        Parameters
+        ----------
+        distance : float
+            The distance to test in metres.
+
+        Returns
+        -------
+        : bool
+            true if the point is within the rayleigh distance false 
+            otherwise.
+        """
+        return np.abs(self.get_position() + distance - \
+            self.location_of_waist()) <= self.rayleigh_distance()
+
+
+    def set_phase_radius(self : GaussianWavefront, 
+            phase_radius : float) -> GaussianWavefront:
+        """
+        Mutator for the phase_radius.
+
+        Parameters
+        ----------
+        phase_radius : float
+            The new phase_radius in meters.
+
+        Returns
+        -------
+        : GaussianWavefront
+            A modified GaussianWavefront with the new phase_radius.
+        """
+        return eqx.tree_at(lambda wavefront : wavefront.phase_radius, 
+            self, phase_radius, is_leaf = lambda leaf : leaf is None)
+
+
+    def set_beam_radius(self : GaussianWavefront, 
+            beam_radius : float) -> GaussianWavefront:
+        """
+        Mutator for the `beam_radius`.
+
+        Parameters
+        ----------
+        beam_radius : float
+            The new beam_radius in meters.
+
+        Returns
+        -------
+        : GaussianWavefront
+            A modified GaussianWavefront with the new beam_radius.
+        """
+        return eqx.tree_at(
+            lambda wavefront : wavefront.beam_radius, self, beam_radius,
+            is_leaf = lambda leaf : leaf is None)


Test patch:
diff --git a/.github/workflows/automated_tests.yml b/.github/workflows/automated_tests.yml
new file mode 100644
index 00000000..5f8f1796
--- /dev/null
+++ b/.github/workflows/automated_tests.yml
@@ -0,0 +1,34 @@
+# Author : Jordan Dennis
+# Date : 29/06/2022
+# Title : automated_tests.yml
+# Description : Automatically runs tests when branches are merged.
+#   The tests that are run can be found in the tests directory.
+
+name: automated tests
+on: [push]
+
+jobs:
+    build:
+        name: automated tests
+        runs-on: ubuntu-20.04
+        steps:
+            - name: checkout
+              uses: actions/checkout@master
+
+            - name: python
+              uses: actions/setup-python@v1
+              with:
+                python-version: 3.10.5
+                architecture: x64
+
+            - name: install environment
+              run: |
+                python -m pip install --upgrade pip
+                pip install --quiet pytest
+                pip install --quiet -r requirements.txt
+                pip install --quiet .
+
+            - name: run tests
+              run: |
+                pytest tests/wavefronts.py
+            
diff --git a/tests/utilities.py b/tests/utilities.py
new file mode 100644
index 00000000..7b178f3c
--- /dev/null
+++ b/tests/utilities.py
@@ -0,0 +1,454 @@
+"""
+utilities.py
+This file contains testing utilities to help with the generation of 
+safe test wavefronts, propagators, layers and detectors. The defined
+classes are:
+"""
+__author__ = "Jordan Dennis"
+__date__ = "28/06/2022"
+
+
+import jax.numpy as numpy
+import dLux
+import typing
+
+
+Array = typing.NewType("Array", numpy.ndarray)
+
+dLuxModule = typing.NewType("dLuxModule", object)
+Wavefront = typing.NewType("Wavefront", object)
+PhysicalWavefront = typing.NewType("PhysicalWavefront", Wavefront)
+AngularWavefront = typing.NewType("AngularWavefront", Wavefront)
+GaussianWavefront = typing.NewType("GaussianWavefront", Wavefront)
+
+Utility = typing.NewType("Utility", object)
+WavefrontUtility = typing.NewType("WavefrontUtility", Utility)
+PhysicalWavefrontUtility = typing.NewType("PhysicalWavefrontUtility", 
+    WavefrontUtility)
+AngularWavefrontUtility = typing.NewType("AngularWavefrontUtility", 
+    WavefrontUtility)
+GaussianWavefrontUtility = typing.NewType("GaussianWavefrontUtility",
+    WavefrontUtility)
+
+UtilityUser = typing.NewType("UtilityUser", object)
+
+
+class UtilityUser(object):
+    """
+    The base utility class. These utility classes are designed to 
+    define safe constructors and constants for testing. These   
+    classes are for testing purposes only. 
+    """
+    utility : Utility
+
+
+    def get_utility(self : UtilityUser) -> Utility:
+        """
+        Accessor for the utility. 
+
+        Returns 
+        -------
+        : Utility
+            The utility
+        """
+        return self.utility 
+
+
+class Utility(object):
+    """
+    """
+
+
+    def __init__(self : Utility) -> Utility:
+        """
+        Construct a new Utility.
+
+        Returns
+        : Utility 
+            The utility. 
+        """
+        raise TypeError("Abstract class Utility should" + \
+            "not be directly substantiated.")
+
+    
+    def construct(self : Utility) -> dLuxModule:
+        """
+        Safe constructor for the dLuxModule, associated with 
+        this utility.
+
+        Returns
+        -------
+        : dLuxModule
+            A safe dLuxModule for testing.
+        """
+        raise TypeError("The abstract construct method" + \
+            "should never be directly invoked")
+
+
+    def approx(self : Utility, result : Array, comparator : Array) -> Array:
+        """
+        Compare two arrays to within floating point precision.
+
+        Parameters
+        ----------
+        result : Array
+            The result that you want to test for nearness to
+            comparator.
+        comparator : Array
+            The comparison array.
+
+        Returns
+        -------
+        : Array[bool]
+            True if the array elements are similar to float 
+            error, False otherwise. 
+        """
+        lower_bound = (result - 0.0005) <= comparator
+        upper_bound = (result + 0.0005) >= comparator
+        return lower_bound & upper_bound 
+
+
+class WavefrontUtility(Utility):
+    """
+    Defines safe state constants and a simple constructor for a safe
+    `Wavefront` object. 
+
+    Attributes
+    ----------
+    offset : Array[float]
+        A simple array defining the angular displacement of the 
+        wavefront. 
+    wavelength : float
+        A safe wavelength for the testing wavefronts in meters
+    """
+    wavelength : float
+    offset : Array
+    size : int
+    amplitude : Array 
+    phase : Array
+    pixel_scale : float
+
+
+    def __init__(self : WavefrontUtility, /, 
+            wavelength : float = None, 
+            offset : Array = None,
+            size : int = None,
+            amplitude : Array = None,
+            phase : Array = None,
+            pixel_scale : float = None) -> WavefrontUtility:
+        """
+        Parameters
+        ----------
+        wavelength : float = 550.e-09
+            The safe wavelength for the utility in meters.
+        offset : Array = [0., 0.]
+            The safe offset for the utility in meters.
+        size : int
+            A parameter for defining consistent wavefront pixel arrays 
+            without causing errors.
+        amplitude : Array[float] 
+            A simple array defining electric field amplitudes without 
+            causing errors.
+        phase : Array[float]
+            A simple array defining the pixel phase for a wavefront, 
+            defined to be safe. 
+
+        Returns 
+        -------
+        : WavefrontUtility 
+            The new utility for generating test cases.
+        """
+        self.wavelength = 550.e-09 if not wavelength else wavelength
+        self.offset = numpy.array([0., 0.]).astype(float) if not \
+            offset else numpy.array(offset).astype(float)           
+        self.size = 128 if not size else size
+        self.amplitude = numpy.ones((self.size, self.size)) if not \
+            amplitude else amplitude
+        self.phase = numpy.zeros((self.size, self.size)) if not \
+            phase else phase
+        self.pixel_scale = 1. if not pixel_scale else pixel_scale
+
+        assert self.size == self.amplitude.shape[0]
+        assert self.size == self.amplitude.shape[1]
+        assert self.size == self.phase.shape[0]
+        assert self.size == self.phase.shape[1]          
+ 
+
+    def construct(self : WavefrontUtility) -> Wavefront:
+        """
+        Build a safe wavefront for testing.
+
+        Returns 
+        -------
+        : Wavefront
+            The safe testing wavefront.
+        """
+        return dLux\
+            .Wavefront(self.wavelength, self.offset)\
+            .update_phasor(self.amplitude, self.phase)\
+            .set_pixel_scale(self.pixel_scale)
+
+
+    def get_wavelength(self : WavefrontUtility) -> float:
+        """
+        Accessor for the wavelength associated with this utility.
+
+        Returns
+        -------
+        : int
+            The wavelength of the utility and hence any `Wavefront` 
+            objects it creates in meters.
+        """
+        return self.wavelength
+
+
+    def get_size(self : WavefrontUtility) -> int:
+        """
+        Accessor for the `size` constant.
+
+        Returns
+        -------
+        : int
+            The side length of a pixel array currently stored.
+        """
+        return self.size
+
+
+    def get_amplitude(self : WavefrontUtility) -> Array:
+        """
+        Accessor for the `amplitude` constant.
+
+        Returns 
+        -------
+        : Array
+            The square array of pixel amplitudes in SI units of 
+            electric field.
+        """
+        return self.amplitude
+
+
+    def get_phase(self : WavefrontUtility) -> Array:
+        """
+        Accessor for the `phase` constant.
+
+        Returns
+        -------
+        : Array
+            The square array of pixel phases in radians.
+        """
+        return self.phase
+
+
+    def get_offset(self : WavefrontUtility) -> Array:
+        """
+        Accessor for the `offset` constant.
+
+        Returns
+        -------
+        : Array
+            The angle that the wavefront makes with the x and 
+            y planes in radians.
+        """
+        return self.offset
+
+
+    def get_pixel_scale(self : WavefrontUtility) -> Array:
+        """
+        Accessor for the `pixel_scale` constant.
+
+        Returns
+        -------
+        : Array
+            The `pixel_scale` associated with the wavefront.
+        """
+        return self.pixel_scale
+
+
+class PhysicalWavefrontUtility(WavefrontUtility):
+    """
+    Defines useful safes state constants as well as a basic 
+    constructor for a safe `PhysicalWavefront`.
+    """
+    def __init__(self : PhysicalWavefrontUtility, /,
+            wavelength : float = None, 
+            offset : Array = None,
+            size : int = None, 
+            amplitude : Array = None, 
+            phase : Array = None) -> PhysicalWavefront:
+        """
+        Parameters
+        ----------
+        wavelength : float 
+            The safe wavelength to use for the constructor in meters.
+        offset : Array[float]
+            The safe offset to use for the constructor in radians.
+        size : int
+            The static size of the pixel arrays.
+        amplitude : Array[float]
+            The electric field amplitudes in SI units for electric
+            field.
+        phase : Array[float]
+            The phases of each pixel in radians. 
+
+        Returns
+        -------
+        : PhysicalWavefrontUtility 
+            A helpful class for implementing the tests. 
+        """
+        super().__init__(wavelength, offset, size, amplitude, phase)
+
+
+    def construct(self : PhysicalWavefrontUtility) -> PhysicalWavefront:
+        """
+        Build a safe wavefront for testing.
+
+        Returns 
+        -------
+        : PhysicalWavefront
+            The safe testing wavefront.
+        """
+        wavefront = dLux\
+            .PhysicalWavefront(self.wavelength, self.offset)\
+            .update_phasor(self.amplitude, self.phase)\
+            .set_pixel_scale(self.pixel_scale)
+
+        return wavefront
+
+
+class AngularWavefrontUtility(WavefrontUtility):
+    """
+    Defines useful safes state constants as well as a basic 
+    constructor for a safe `PhysicalWavefront`.
+    """
+    def __init__(self : AngularWavefrontUtility, /,
+            wavelength : float = None, 
+            offset : Array = None,
+            size : int = None, 
+            amplitude : Array = None, 
+            phase : Array = None) -> AngularWavefrontUtility:
+        """
+        Parameters
+        ----------
+        wavelength : float 
+            The safe wavelength to use for the constructor in meters.
+        offset : Array[float]
+            The safe offset to use for the constructor in radians.
+        size : int
+            The static size of the pixel arrays.
+        amplitude : Array[float]
+            The electric field amplitudes in SI units for electric
+            field.
+        phase : Array[float]
+            The phases of each pixel in radians. 
+
+        Returns
+        -------
+        : AngularWavefrontUtility 
+            A helpful class for implementing the tests. 
+        """
+        super().__init__(wavelength, offset, size, amplitude, phase)
+
+
+    def construct(self : AngularWavefrontUtility) -> AngularWavefront:
+        """
+        Build a safe wavefront for testing.
+
+        Returns 
+        -------
+        : PhysicalWavefront
+            The safe testing wavefront.
+        """
+        wavefront = dLux\
+            .AngularWavefront(self.wavelength, self.offset)\
+            .update_phasor(self.amplitude, self.phase)\
+            .set_pixel_scale(pixel_scale)
+            
+        return wavefront
+
+
+class GaussianWavefrontUtility(PhysicalWavefrontUtility):
+    """
+    Defines safe state constants and a simple constructor for a 
+    safe state `GaussianWavefront` object. 
+
+    Attributes
+    ----------
+    beam_radius : float
+        A safe radius for the GaussianWavefront in meters.
+    phase_radius : float
+        A safe phase radius for the GaussianWavefront in radians.
+    position : float
+        A safe position for the GaussianWavefront in meters.
+    """
+    beam_radius : float 
+    phase_radius : float
+    position : float
+
+
+    def __init__(self : GaussianWavefrontUtility, 
+            wavelength : float = None,
+            offset : Array = None,
+            size : int = None,
+            amplitude : Array = None,
+            phase : Array = None,
+            beam_radius : float = None,
+            phase_radius : float = None,
+            position : float = None) -> GaussianWavefrontUtility:
+        """
+        Parameters
+        ----------
+        wavelength : float 
+            The safe wavelength to use for the constructor in meters.
+        offset : Array[float]
+            The safe offset to use for the constructor in radians.
+        size : int
+            The static size of the pixel arrays.
+        amplitude : Array[float]
+            The electric field amplitudes in SI units for electric
+            field.
+        phase : Array[float]
+            The phases of each pixel in radians.
+        beam_radius : float 
+            The radius of the gaussian beam in meters.
+        phase_radius : float
+            The phase radius of the gaussian beam in radians.
+        position : float
+            The position of the gaussian beam in meters.
+
+        Returns
+        -------
+        : GaussianWavefrontUtility 
+            A helpful class for implementing the tests. 
+        """
+        super().__init__(wavelength, offset, size, amplitude, phase)
+        self.beam_radius = 1. if not beam_radius else beam_radius
+        self.phase_radius = numpy.inf if not phase_radius else phase_radius
+        self.position = 0. if not position else position
+
+
+    # TODO: get_beam_radius and get_position
+    def get_phase_radius(self : GaussianWavefrontUtility) -> float:
+        """
+        """
+        return self.phase_radius
+
+
+    def construct(self : GaussianWavefrontUtility) -> GaussianWavefront:
+        """
+        Build a safe wavefront for testing.
+
+        Returns 
+        -------
+        : PhysicalWavefront
+            The safe testing wavefront.
+        """
+        wavefront = dLux\
+            .GaussianWavefront(self.offset,
+               self.wavelength)\
+            .update_phasor(self.amplitude, self.phase)\
+            .set_pixel_scale(self.pixel_scale)\
+            .set_position(0.)\
+            .set_phase_radius(numpy.inf)\
+            .set_beam_radius(1.)
+
+        return wavefront
diff --git a/tests/wavefronts.py b/tests/wavefronts.py
new file mode 100644
index 00000000..334d6ffd
--- /dev/null
+++ b/tests/wavefronts.py
@@ -0,0 +1,701 @@
+import jax.numpy as numpy
+import pytest
+import jax.random as random
+import typing
+
+from dLux import Wavefront
+from dLux import PhysicalWavefront
+from dLux import AngularWavefront
+from dLux import GaussianWavefront
+
+from utilities import WavefrontUtility
+from utilities import PhysicalWavefrontUtility
+from utilities import GaussianWavefrontUtility
+from utilities import AngularWavefrontUtility
+from utilities import UtilityUser
+
+
+Array = typing.NewType("Array", numpy.ndarray)
+
+TestWavefront = typing.NewType("TestWavefront", UtilityUser)
+TestPhysicalWavefront = typing.NewType("TestPhysicalWavefront", UtilityUser)
+TestAngularWavefront = typing.NewType("TestAngularWavefront", UtilityUser)
+TestGaussianWavefront = typing.NewType("TestGaussianWavefront", UtilityUser)
+
+
+class TestWavefront(UtilityUser):
+    """
+    Tests the Wavefront class. Tests are written looking 
+    for null and nan values, with only simple end to end tests.
+
+    Attributes
+    ----------
+    utility : WavefrontUtility
+        A helper for generating safe test cases.
+    """    
+    utility : WavefrontUtility = WavefrontUtility()
+
+
+    def test_get_wavelength(self : TestWavefront) -> None:
+        """
+        Test for the accessor get_wavelength.
+        """
+        OUTPUT_WAVELENGTH = self\
+            .get_utility()\
+            .get_wavelength()
+
+        wavefront = self\
+            .get_utility()\
+            .construct()
+
+        assert wavefront.get_wavelength() == OUTPUT_WAVELENGTH
+            
+
+
+    def test_set_wavelength(self : TestWavefront) -> None:
+        """
+        Test for the mutator set_wavelength.
+        """
+        OLD_WAVELENGTH = self.get_utility().get_wavelength()
+        NEW_WAVELENGTH = 540.e-09
+
+        wavefront = self\
+            .get_utility()\
+            .construct()\
+            .set_wavelength(NEW_WAVELENGTH)
+
+        assert wavefront.get_wavelength() == NEW_WAVELENGTH
+        assert wavefront.get_wavelength() != OLD_WAVELENGTH
+
+
+    def test_get_offset(self : TestWavefront) -> None:
+        """
+        Test for the accessor method get_offset.
+        """
+        wavefront = self\
+            .get_utility()\
+            .construct()
+
+        is_correct = (wavefront.get_offset() == self\
+            .get_utility()\
+            .get_offset())\
+            .all()
+
+        assert is_correct
+
+
+    def test_set_offset(self : TestWavefront) -> None:
+        """
+        Test for the mutator method set_offset.
+        """
+        NEW_OFFSET = numpy.array([numpy.pi] * 2)
+        OLD_OFFSET = self\
+            .get_utility()\
+            .get_offset()
+
+        wavefront = self\
+            .get_utility()\
+            .construct()\
+            .set_offset(NEW_OFFSET)
+
+        assert (wavefront.get_offset() == NEW_OFFSET).all()
+        assert (wavefront.get_offset() != OLD_OFFSET).all()
+    
+
+    def test_get_real(self : TestWavefront) -> None:
+        """
+        Rotates the phasor through real -> imaginary and checks that 
+        the output is correct.         
+        """
+        FULL_REAL = 0.
+        PART_REAL = numpy.pi / 4
+        NONE_REAL = numpy.pi / 2
+
+        FULL_REAL_OUTPUT = 1.
+        PART_REAL_OUTPUT = 1. / numpy.sqrt(2.)
+        NONE_REAL_OUTPUT = 0.        
+
+        full_real_wavefront = self\
+            .utility\
+            .construct()
+
+        part_real_wavefront = self\
+            .utility\
+            .construct()\
+            .add_phase(PART_REAL)
+
+        none_real_wavefront = self\
+            .utility\
+            .construct()\
+            .add_phase(NONE_REAL)
+
+        is_full_real_correct = self\
+            .get_utility()\
+            .approx(full_real_wavefront.get_real(), FULL_REAL_OUTPUT)\
+            .all()
+ 
+        is_part_real_correct = self\
+            .get_utility()\
+            .approx(part_real_wavefront.get_real(), PART_REAL_OUTPUT)\
+            .all()
+
+        is_none_real_correct = self\
+            .get_utility()\
+            .approx(none_real_wavefront.get_real(), NONE_REAL_OUTPUT)\
+            .all()     
+
+        print(is_full_real_correct)
+
+        assert is_full_real_correct
+        assert is_part_real_correct
+        assert is_none_real_correct    
+        
+
+    def test_get_imaginary(self : TestWavefront) -> None:
+        """
+        Rotates the phasor through real -> imaginary and checks that 
+        the output is correct
+        """
+        NONE_IMAGINARY = numpy.pi
+        PART_IMAGINARY = numpy.pi / 4
+        FULL_IMAGINARY = numpy.pi / 2
+
+        NONE_IMAGINARY_OUT = 0.
+        PART_IMAGINARY_OUT = 1. / numpy.sqrt(2)
+        FULL_IMAGINARY_OUT = 1.
+
+        none_imaginary_wavefront = self\
+            .utility\
+            .construct()\
+            .add_phase(NONE_IMAGINARY)
+
+        part_imaginary_wavefront = self\
+            .utility\
+            .construct()\
+            .add_phase(PART_IMAGINARY)
+
+        full_imaginary_wavefront = self\
+            .utility\
+            .construct()\
+            .add_phase(FULL_IMAGINARY)
+
+        assert self\
+            .get_utility()\
+            .approx(full_imaginary_wavefront.get_imaginary(), 
+                FULL_IMAGINARY_OUT)\
+            .all()
+
+        assert self\
+            .get_utility()\
+            .approx(part_imaginary_wavefront.get_imaginary(),
+                PART_IMAGINARY_OUT)\
+            .all()
+  
+        assert self\
+            .get_utility()\
+            .approx(none_imaginary_wavefront.get_imaginary(),
+                NONE_IMAGINARY_OUT)\
+            .all() 
+
+
+    def test_multiply_amplitude(self : TestWavefront) -> None:
+        """
+        Checks that the amplitude array has been updated, after 
+        operations
+        """
+        INITIAL_AMPLITUDE = 1.
+        CHANGED_AMPLITUDE = 2.
+
+        initial_wavefront = self.get_utility().construct()
+        initial_amplitude = initial_wavefront.get_amplitude()
+        changed_wavefront = initial_wavefront.multiply_amplitude(2)
+        changed_amplitude = changed_wavefront.get_amplitude()
+        
+        assert (initial_amplitude == INITIAL_AMPLITUDE).all()  
+        assert (changed_amplitude == CHANGED_AMPLITUDE).all()  
+        
+
+    def test_add_phase(self : TestWavefront) -> None:   
+        """
+        Checks that the phase array is correctly updated by the 
+        operations. 
+        """
+        INITIAL_PHASE = 0.
+        CHANGED_PHASE = numpy.pi
+
+        initial_wavefront = self.get_utility().construct()
+        initial_phase = initial_wavefront.get_phase()
+        changed_wavefront = initial_wavefront.add_phase(numpy.pi)
+        changed_phase = changed_wavefront.get_phase()
+        
+        assert (initial_phase == INITIAL_PHASE).all()
+        assert (changed_phase == CHANGED_PHASE).all()
+
+
+    def test_update_phasor(self : TestWavefront) -> None:
+        """
+        Checks that the phasor is correctly updated by changes to
+        both the phase and the amplitude. 
+        """
+        NEW_AMPLITUDE = numpy.ones(
+            (self.utility.size, self.utility.size))
+        NEW_PHASE = numpy.ones(
+            (self.utility.size, self.utility.size))
+
+        wavefront = self\
+            .get_utility()\
+            .construct()\
+            .update_phasor(NEW_AMPLITUDE, NEW_PHASE)
+
+        assert (wavefront.get_phase() == NEW_PHASE).all() 
+        assert (wavefront.get_amplitude() == NEW_AMPLITUDE).all() 
+        
+
+    def test_wavefront_to_point_spread_function(
+            self : TestWavefront) -> None:
+        """
+        Test that the point spread function is correctly generated
+        from the amplitude and phase arrays. Considered correct
+        if the output is the amplitude ** 2 and modifying the phase 
+        does not affect the PSF
+        """
+        OUTPUT_AMPLITUDE = self.utility.amplitude * 2.
+        OUTPUT_PSF = OUTPUT_AMPLITUDE ** 2
+
+        output_psf = self\
+            .get_utility()\
+            .construct()\
+            .set_amplitude(OUTPUT_AMPLITUDE)\
+            .wavefront_to_psf()
+         
+        assert (output_psf == OUTPUT_PSF).all() 
+
+
+    def test_add_optical_path_difference(
+            self : TestWavefront) -> None:
+        """
+        Testing for correct behaviour when optical paths length 
+        distances are added. 
+        """
+        INITIAL_PHASE = 0.
+        CHANGED_PHASE = numpy.pi
+        WAVELENGTH = self.get_utility().get_wavelength()
+
+        initial_wavefront = self.get_utility().construct()
+        initial_phase = initial_wavefront.get_phase()
+        changed_wavefront = initial_wavefront.add_opd(WAVELENGTH / 2)
+        changed_phase = changed_wavefront.get_phase()
+
+        assert (initial_phase == INITIAL_PHASE).all() 
+        assert (changed_phase == CHANGED_PHASE).all() 
+
+
+    def test_normalise(self : TestWavefront) -> None:
+        """
+        Checks that the normalise functionality is working. The 
+        behaviour is considered functional if the maximum 
+        value encountered is 1. and the minimum value encountered 
+        is 0.
+        """
+        key = random.PRNGKey(0)
+        size = self.get_utility().get_size() 
+        
+        INITIAL_AMPLITUDE = random.normal(key, (size, size))
+        OUTPUT = INITIAL_AMPLITUDE / numpy.linalg.norm(INITIAL_AMPLITUDE)
+
+        normalised_amplitude = self\
+            .get_utility()\
+            .construct()\
+            .set_amplitude(INITIAL_AMPLITUDE)\
+            .normalise()\
+            .get_amplitude()
+
+        is_correct = self\
+            .get_utility()\
+            .approx(normalised_amplitude, OUTPUT)\
+            .all()
+
+        assert is_correct 
+        
+
+    def test_get_pixel_position_vector(
+            self : TestWavefront) -> None:
+        """
+        The get_xs_vec() function is consider correct if the 
+        minimum value -(npix - 1) / 2 and the maximum value is
+        (npix - 1) / 2
+        """
+        size = self\
+            .get_utility()\
+            .get_size()
+
+        pixel_coordinates = self\
+            .get_utility()\
+            .construct()\
+            .get_pixel_coordinates(size)
+        
+        assert (pixel_coordinates.max() == (size - 1) / 2).all()
+        assert (pixel_coordinates.min() == -(size - 1) / 2).all()
+
+
+    def test_get_pixel_grid(self : TestWavefront) -> None:
+        """
+        The get_XXYY function is considered correct if it produces 
+        an array that is of dimensions (2, self.SIZE, self.SIZE)
+        as it uses the get_xs_vec() function under the hood.
+        """
+        size = self\
+            .get_utility()\
+            .get_size()
+
+        pixel_grid = self\
+            .get_utility()\
+            .construct()\
+            .get_pixel_grid()
+        
+        assert pixel_grid.shape == (2, size, size)
+
+
+    # TODO: Update the docstrings to match the new implementations
+    def test_get_pixel_positions(self : TestWavefront) -> None:
+        """
+        The get_xycoords() function is considered correct if it 
+        produces an array with the correct dimesions such that the 
+        minimum and maximum values are plus and minus
+        self.pixelscale * (npix - 1) / 2
+        """
+        size = self\
+            .get_utility()\
+            .get_size()
+
+        pixel_scale = self\
+            .get_utility()\
+            .get_pixel_scale()
+
+        physical_coordinates = self\
+            .get_utility()\
+            .construct()\
+            .get_pixel_positions()
+        
+        assert physical_coordinates.max() == (size - 1) / 2 * pixel_scale
+        assert physical_coordinates.min() == -(size - 1) / 2 * pixel_scale 
+        assert physical_coordinates.shape == (2, size, size)
+
+
+    def test_set_phase(self : TestWavefront) -> None:
+        """
+        Functionality that is not currently supported, allows the 
+        state to be immutably changed and viewed from outside the 
+        namespace. 
+        """
+        size = self.get_utility().get_size()
+        key = random.PRNGKey(0)
+        phase = random.normal(key, (size, size))
+
+        wavefront = self\
+            .get_utility()\
+            .construct()\
+            .set_phase(phase)
+        
+        assert (wavefront.get_phase() == phase).all() 
+ 
+
+    def test_set_amplitude(self : TestWavefront) -> None:
+        """
+        Functionality that is not currently supported, provides 
+        immutable access to the state of the wavefront. 
+        """
+        size = self.get_utility().get_size()
+        key = random.PRNGKey(0)
+        amplitude = random.normal(key, (size, size))
+
+        wavefront = self\
+            .get_utility()\
+            .construct()\
+            .set_amplitude(amplitude)
+        
+        assert (wavefront.get_amplitude() == amplitude).all() 
+
+
+class TestPhysicalWavefront(UtilityUser):
+    """
+    Tests the dLux.PhysicalWavefront class. 
+
+    Attributes
+    ----------
+    utility : PhysicalWavefrontUtility
+        A helper object for generating safe test cases.
+    """
+    utility : PhysicalWavefrontUtility = PhysicalWavefrontUtility()
+
+
+    def test_constructor(self : TestPhysicalWavefront) -> None:
+        """
+        Tests the `PhysicaleWavefront` constructor.
+        """
+        WAVELENGTH = self.get_utility().get_wavelength()
+        OFFSET = self.get_utility().get_offset()
+
+        wavefront = PhysicalWavefront(WAVELENGTH, OFFSET)
+        
+        assert wavefront.get_wavelength() == WAVELENGTH
+        assert (wavefront.get_offset() == OFFSET).all()
+        assert wavefront.get_amplitude() == None
+        assert wavefront.get_phase() == None
+        
+
+
+class TestAngularWavefront(UtilityUser):
+    """
+    Tests the `AngularWavefront` class. As `PhysicalWavefront` is 
+    just a type ecosystem extension this only tests the constructor. 
+
+    Attributes 
+    ----------
+    utility : AngularWavefrontUtility
+        A helper class for generating safe test cases.
+    """
+    utility : AngularWavefrontUtility = AngularWavefrontUtility()
+
+
+    def test_constructor(self : TestAngularWavefront) -> None:
+        """
+        Checks that the constructor initialises all the required 
+        fields.
+        """
+        WAVELENGTH = self.get_utility().get_wavelength()
+        OFFSET = self.get_utility().get_offset()
+
+        wavefront = AngularWavefront(WAVELENGTH, OFFSET)
+        
+        assert wavefront.get_wavelength() == WAVELENGTH
+        assert (wavefront.get_offset() == OFFSET).all()
+        assert wavefront.get_amplitude() == None
+        assert wavefront.get_phase() == None
+
+
+class TestGaussianWavefront(UtilityUser):
+    """
+    Tests the exteded functionality of the `GaussianWavefront`
+    class.
+
+    Attributes
+    ----------
+    utility : GaussianWavefrontUtility
+        A helper class for testing the `GaussianWavefront`.
+    """
+    utility : GaussianWavefront = GaussianWavefrontUtility()
+
+
+    def test_set_position(self : TestGaussianWavefront) -> None:
+        """
+        Provides immutable access to the state of the wavefront.
+        Considered cocrect if the initial instance is not modified 
+        and a modified wavefront is created.
+        """
+        initial_wavefront = self.get_utility().construct()
+        changed_wavefront = initial_wavefront.set_position(1.)
+
+        assert (initial_wavefront.get_position() == 0.)
+        assert (changed_wavefront.get_position() == 1.)
+
+
+    def test_set_phase_radius(self : TestGaussianWavefront) -> None:
+        """
+        Provides immutable access to the state of the wavefront. 
+        Considered correct if the initial instance is not modified
+        and a modified wavefront is created.
+        """
+        initial_wavefront = self.get_utility().construct()
+        changed_wavefront = initial_wavefront.set_phase_radius(1.)
+
+        # TODO: Work out a good default value for this. 
+        assert initial_wavefront.get_phase_radius() == self\
+            .get_utility()\
+            .get_phase_radius()
+
+        assert changed_wavefront.get_phase_radius() == 1.
+       
+
+    # Properties 
+    def test_rayleigh_distance(self : TestGaussianWavefront) -> None:
+        """
+        Checks that the rayleigh distance is calculated correctly
+        based on the parameters that are input. 
+        """
+        wavefront = self.get_utility().construct()
+        rayleigh_distance = wavefront.rayleigh_distance()
+        correct_rayleigh_distance = numpy.pi * \
+            wavefront.get_beam_radius() ** 2 / \
+            wavefront.get_wavelength()        
+
+        assert rayleigh_distance == correct_rayleigh_distance
+               
+
+    def test_location_of_waist(self : TestGaussianWavefront) -> None:
+        """
+        Checks that the location of the waist is correctly determined.
+        Simply runs the explicit calculations externally and compares.
+        """
+        wavefront = self\
+            .get_utility()\
+            .construct()\
+            .set_phase_radius(1.)
+
+        correct_location_of_waist = - wavefront.get_phase_radius() / \
+            (1 + (wavefront.get_phase_radius() / \
+                wavefront.rayleigh_distance()) ** 2)
+
+        assert (wavefront.location_of_waist() == correct_location_of_waist)
+
+
+    def test_waist_radius(self : TestGaussianWavefront) -> None:
+        """
+        Directly confirms that the correct numerical calculations are 
+        implemented, by the method
+        """
+        wavefront = self.get_utility().construct()
+        correct_waist_radius = wavefront.get_beam_radius() / \
+            numpy.sqrt(1 + (wavefront.rayleigh_distance() / \
+            wavefront.get_beam_radius()) ** 2)
+
+        assert wavefront.waist_radius() == correct_waist_radius        
+
+
+    def test_transfer_function_not_nan(
+            self : TestGaussianWavefront) -> None:
+        """
+        Check the boundary case distance = 0. and then two normal
+        inputs a large and a small.
+        """
+        wavefront = self.get_utility().construct()
+        zero = wavefront.transfer_function(0.)
+        small = wavefront.transfer_function(0.01)
+        large = wavefront.transfer_function(1.)
+
+        # assert not numpy.isnan(zero).any()
+        assert not numpy.isnan(small).any()
+        assert not numpy.isnan(large).any()
+
+
+    def test_transfer_function_not_inf(
+            self : TestGaussianWavefront) -> None:
+        """
+        Checks that the boundary case distance == 0. and then two 
+        normal inputs do not produce infinte results
+        """
+        wavefront = self.get_utility().construct()
+        zero = wavefront.transfer_function(0.)
+        small = wavefront.transfer_function(0.01)
+        large = wavefront.transfer_function(1.)
+
+        # assert not numpy.isinf(zero).any()
+        assert not numpy.isinf(small).any()
+        assert not numpy.isinf(large).any()
+
+
+    def test_quadratic_phase_factor_not_nan(
+            self : TestGaussianWavefront) -> None:
+        """
+        Checks the boundary case distance == 0. for nan inputs 
+        as well as a small and a large typical use case
+        """
+        wavefront = self.get_utility().construct()
+        # zero = wavefront.quadratic_phase_factor(0.)
+        # infinte = wavefront.quadratic_phase_factor(numpy.inf)
+        small = wavefront.quadratic_phase_factor(0.001)
+        large = wavefront.quadratic_phase_factor(1.)
+
+        # assert not numpy.isnan(zero).any()
+        # assert not numpy.isnan(infinte).any()
+        assert not numpy.isnan(small).any()
+        assert not numpy.isnan(large).any()        
+
+
+    def test_quadratic_phase_factor_not_inf(
+            self : TestGaussianWavefront) -> None:
+        """
+        Tests the boundary cases distance == numpy.inf and distance 
+        == 0. as well as a small and a large typical case for 
+        infinite values
+        """
+        wavefront = self.get_utility().construct()
+        # zero = wavefront.quadratic_phase_factor(0.)
+        # infinte = wavefront.quadratic_phase_factor(numpy.inf)
+        small = wavefront.quadratic_phase_factor(0.001)
+        large = wavefront.quadratic_phase_factor(1.)
+
+        # assert not numpy.isinf(zero).any()
+        # assert not numpy.isinf(infinte).any()
+        assert not numpy.isinf(small).any()
+        assert not numpy.isinf(large).any()  
+
+        
+    def test_pixel_scale_not_nan(
+            self : TestGaussianWavefront) -> None:
+        """
+        Checks that the new pixel scale is not generated to be nan 
+        by a negative, zero and positive use case
+        """
+        wavefront = self.get_utility().construct()
+        negative = wavefront.calculate_pixel_scale(-0.01)
+        zero = wavefront.calculate_pixel_scale(0.)
+        positive = wavefront.calculate_pixel_scale(0.01)
+
+        assert not numpy.isnan(negative).any()
+        assert not numpy.isnan(zero).any()
+        assert not numpy.isnan(positive).any()        
+
+
+    def test_pixel_scale_not_inf(
+            self : TestGaussianWavefront) -> None:
+        """
+        Tests the pixel scale for infinite values when passed a 
+        negative, zero and positive position.
+        """
+        wavefront = self.get_utility().construct()
+        negative = wavefront.calculate_pixel_scale(-0.01)
+        zero = wavefront.calculate_pixel_scale(0.)
+        positive = wavefront.calculate_pixel_scale(0.01)
+
+        assert not numpy.isinf(negative).any()
+        assert not numpy.isinf(zero).any()
+        assert not numpy.isinf(positive).any()       
+
+
+    def test_is_inside(self : TestGaussianWavefront) -> None:
+        """
+        Branch coverage for a two dimensional is_inside call. That is 
+        to say:
+        
+        false, false
+        false, true
+        true, false
+        true, true
+        """
+        wavefront = self\
+            .get_utility()\
+            .construct()\
+            .set_phase_radius(1.)
+
+        rayleigh_distance = wavefront.rayleigh_distance()
+
+        false_false = wavefront.is_inside(
+            rayleigh_distance * numpy.ones((2, )) + 1.)
+
+        false_true = wavefront.is_inside(numpy\
+            .zeros((2, ))\
+            .at[0]\
+            .set(numpy.squeeze(rayleigh_distance) + 1.))
+
+        true_false = wavefront.is_inside(numpy\
+            .zeros((2, ))\
+            .at[1]\
+            .set(numpy.squeeze(rayleigh_distance) + 1.))
+
+        true_true = wavefront.is_inside(numpy.array([0., 0.]))
+
+        assert (false_false == numpy.array([False, False])).all() 
+        assert (false_true == numpy.array([False, True])).all() 
+        assert (true_false == numpy.array([True, False])).all() 
+        assert (true_true == numpy.array([True, True])).all()