Repository: LouisDesdoigts/dLux (language: python)

Issue:
Restrucutre /src
So the /src folder is getting a bit unwieldy. I propose we create a sub-directory called `extras` or something to store the non-core functionality scripts.

The idea would be the scripts recently moved into /src such as `bayes.py`, `helpers.py`, `plotting.py` would live in /extras as a collection of scripts that are designed to work with ∂lux but are not core and therefore not documented or tested.

Over time some of this functionality would be lifted to /src, such as a subset of the `helpers.py` functions as they are widely applicable to the usage of the package.

This could also be a good place to put a `development.py` script to place layers in development (ie not documented and tested) so that we can all have access and use consistent versions. This would be very useful for the development of the JWST models so we can put code to share without while still being flexible.
`Image` Class to Reduce repetition
Hi all,
A large number of classes store pixel arrays a.k.a images. There is certain functionality that is associated with this, for example the generation of para-axial arrays. I have tried to keep the naming consistent but we have `Propagator._get_pixel_grid` and `Wavefront.get_pixel_positions` which do the same thing. There is the option to make `Wavefront.get_pixel_positions` an `@staticmethod` with `npix` as a parameter. 

Aside: A static method does not require the class to instantiate an object and can literally be called like `Wavefront.get_pixel_positions(npix)`.

While this is probably the neatest solution at present it does not make much physical sense. Why should the wavefront describe that behaviour? However, I do think we should include the logic of the function in only one place. I think this because `Propagator` and `Wavefront` actually defined the para-axial centre differently and when I discussed this with @LouisDesdoigts we decided that it was a bug. Additionally, some `Layer` classes, use this functionality, for example the `CircularApperture`. 

The typical OOP solution is to build a class `Image` or `ImageHolder` that the other classes can inherit from (thankfully python supports multiple inheritance). This would hold the parameter `_image : Array` and define a series of useful methods. 
```py
class Image():
    _image : Array
    
    # NOTE: names are not final, just suggestions
    def get_pixel_coordinates() -> Vector: # Along edge
    def get_pixel_positions() -> Tensor: # 2 stacked grids in 3rd order tensor 
    def get_radial_positions() -> Matrix: # stacked grid
```
The downfall of this is that you get a very complicated class hierarchy (still quite simple in our case), which is the general downfall of OOP. 

A further option is to register all the methods of the class I describe above as `@staticmethods` making the nearest thing to a static class python provides. With reference to #58 this could go into `src/utilities/image.py` and then usage would be through `Image.get_pixel_coordinates(npix)` ect. 

I quite like the idea of implementing such a static class, because it also avoids subclasses ending up with lots of methods that they never use as can happen with OOP programming. 

As always let me know what you think.

Regards

Jordan
Restrucutre /src
So the /src folder is getting a bit unwieldy. I propose we create a sub-directory called `extras` or something to store the non-core functionality scripts.

The idea would be the scripts recently moved into /src such as `bayes.py`, `helpers.py`, `plotting.py` would live in /extras as a collection of scripts that are designed to work with ∂lux but are not core and therefore not documented or tested.

Over time some of this functionality would be lifted to /src, such as a subset of the `helpers.py` functions as they are widely applicable to the usage of the package.

This could also be a good place to put a `development.py` script to place layers in development (ie not documented and tested) so that we can all have access and use consistent versions. This would be very useful for the development of the JWST models so we can put code to share without while still being flexible.
Development of optical components layer script
So we want to support arbitrary optical systems configurations. This would mean constructing a set of very general layers for common optical components, ie lenses, phase plates, etc. @Jordan-Dennis has made a start on a structure for this script which looks rather good. In the long run we could build a minimal web-scraper to map specific optical components like [this](https://www.thorlabs.com/newgrouppage9.cfm?objectgroup_id=12767) to their exact chromatic response and manufacturer specifications. 

Integration with manufacturers is probably a unnecessary long-term goal, but the basic optical components script is desired.
Development of optical components layer script
So we want to support arbitrary optical systems configurations. This would mean constructing a set of very general layers for common optical components, ie lenses, phase plates, etc. @Jordan-Dennis has made a start on a structure for this script which looks rather good. In the long run we could build a minimal web-scraper to map specific optical components like [this](https://www.thorlabs.com/newgrouppage9.cfm?objectgroup_id=12767) to their exact chromatic response and manufacturer specifications. 

Integration with manufacturers is probably a unnecessary long-term goal, but the basic optical components script is desired.
Development of optical components layer script
So we want to support arbitrary optical systems configurations. This would mean constructing a set of very general layers for common optical components, ie lenses, phase plates, etc. @Jordan-Dennis has made a start on a structure for this script which looks rather good. In the long run we could build a minimal web-scraper to map specific optical components like [this](https://www.thorlabs.com/newgrouppage9.cfm?objectgroup_id=12767) to their exact chromatic response and manufacturer specifications. 

Integration with manufacturers is probably a unnecessary long-term goal, but the basic optical components script is desired.
Hexagons and hexikes
Hi all,
I have not allocated a position to the new `layers.py` which contains __only__ the new `HexagonalBasis` code and friends. This code has been visually tested very thoroughly against `poppy` although I have held of on writing formal unit tests on it because I _know_ it works and there are more pressing things to develop. 

The code is `jit`-able but due to the overhead incurred by the additional vectorizations I have made it is actually slower than the `numpy` version even after `jit`. This should be fine because we are just saving to a file with the `save(filename : str)` method. 

I have not introduced a `CompoundAperture` class to deal with saving multiple basis in tensors to files. RSVP me if we want to add this functionality. Currently we can rotate, shear (spelled correctly :laughing:), translate and magnify the basis vectors. 

As a final note; conceding to `jax` we have our first constant `MAX_DIFF`, if you want to generate more the 21 basis vectors you will have to increase this to 6 or above, but I doubt we will. For lower memory use and slightly better speed you can drop it to 2 if you are only generating the first nine. 

@LouisDesdoigts, if you can review the code after you have settled in :canada:, and leave a note allocating the position that you want me to put the code and I will move it there; I was thinking maybe a `basis.py` in the `dLux/dLux` directory.  

Regards 

Jordan


Discussion / hints:
I agree, we should implement some set of sub-modules. Personally I quite like the following structure:
```
src/
    optics/
        wavefronts.py
        layers.py
        base.py
    statistics/
        bayes.py
        helpers.py?
    graphics/
        plotting.py
```
As this groups things by the functionality that they add. We would have to change the init file to:
```
from . import optics
from . import statistics
from . import graphics
```
So the import in scripts would becomes `from dLux.optics import *`. 
    

I agree, we should implement some set of sub-modules. Personally I quite like the following structure:
```
src/
    optics/
        wavefronts.py
        layers.py
        base.py
    statistics/
        bayes.py
        helpers.py?
    graphics/
        plotting.py
```
As this groups things by the functionality that they add. We would have to change the init file to:
```
from . import optics
from . import statistics
from . import graphics
```
So the import in scripts would becomes `from dLux.optics import *`. 
    
So my initial thoughts are included below. This is what @LouisDesdoigts and I originally discussed:
```py
class Layer(eqx.Module):
    def __call__(self, parameters): # abstract

class CurvedOptic(Layer):
    def _curvature(self): # abstract

class Mirror(Layer):
    def _reflect(self, wavefront): # abstract
    def __call__(self, parameters): # concrete

class GaussianMirror(CurvedOptic, Mirror):
    def _curvature(self): # concrete
    def _reflect(self, wavefront): # concrete

class HyperbolicMirror(CurvedOptic, Mirror):
    def _curvature(self): # concrete
    def _reflect(self, wavefront): # concrete

class SphericalMirror(CurvedOptic, Mirror):
    def _curvature(self): # concrete
    def _reflect(self, wavefront): # concrete

class Lens(CurvedOptic):
    index_of_refraction : float 
    def _curvature(self): # abstract
    def _focus(self, wavefront): # concrete
    def __call__(self, parameters): # concrete

class GaussianLens(Lens):
    def _curvature(self): # concrete

class HyperbolicLens(Lens):
    def _curvature(self): # concrete 

class ParabolicLens(Lens):
    def _curvature(self): # concrete

class Aperture(Optic):
    def _area(self): # Still abtract 
    def __call__(self): # concrete

class CircularAperture(Aperture):
    radius : float
    def _area(self, wavefront):

class RegularPolygonalAperture(Aperture):
    number_of_sides : int
    radius : float
    def _area(self): # concrete

class ArbitraryAperture(Aperture):
    binary_mask : int
    def _area(self):
        return self.binary_mask

class PhasePlate(Optic):
    def _phase_shift(self): # abstract
    def __call__(self): # concrete

class ArbitraryPhasePlate(Optic):
    phase : array
    def _phase_shift(self):
        return self.phase

class ChromaticPhasePlate(Optic):
    def _phase_shift(self, wavefront): # concrete
        # dependence on wavelength 
```
Although I was thinking about #60 and thought we could extend this even further by adding:
```py
class Source(Layer): # abstract
    def _create_wavefront(self): # abstract
    def __call__(self, parameters): # concrete

class CoherentPointSource(Layer):
    def _create_wavefront(self): # concrete

class CoherentPlaneSource(Layer): # Point source at infinty?
    def _create_wavefront(self): # concrete
```
where there are many other sources that could be generated. This leads to the natural extension of `CreateWavefront` into a set of many simple layers with specific properties. To me, it makes most sense to have `CoherentPointSource` output an `AngularWavefront` since the `Wavefront` is spherical and is best described by a field of view. `CoherentPlaneSource` on the other hand is better suited to a `PhysicalWavefront`. `GaussianWavefront` is a bit trickier because it requires a `Lens` or `Mirror` with a curvature. 

So my initial thoughts are included below. This is what @LouisDesdoigts and I originally discussed:
```py
class Layer(eqx.Module):
    def __call__(self, parameters): # abstract

class CurvedOptic(Layer):
    def _curvature(self): # abstract

class Mirror(Layer):
    def _reflect(self, wavefront): # abstract
    def __call__(self, parameters): # concrete

class GaussianMirror(CurvedOptic, Mirror):
    def _curvature(self): # concrete
    def _reflect(self, wavefront): # concrete

class HyperbolicMirror(CurvedOptic, Mirror):
    def _curvature(self): # concrete
    def _reflect(self, wavefront): # concrete

class SphericalMirror(CurvedOptic, Mirror):
    def _curvature(self): # concrete
    def _reflect(self, wavefront): # concrete

class Lens(CurvedOptic):
    index_of_refraction : float 
    def _curvature(self): # abstract
    def _focus(self, wavefront): # concrete
    def __call__(self, parameters): # concrete

class GaussianLens(Lens):
    def _curvature(self): # concrete

class HyperbolicLens(Lens):
    def _curvature(self): # concrete 

class ParabolicLens(Lens):
    def _curvature(self): # concrete

class Aperture(Optic):
    def _area(self): # Still abtract 
    def __call__(self): # concrete

class CircularAperture(Aperture):
    radius : float
    def _area(self, wavefront):

class RegularPolygonalAperture(Aperture):
    number_of_sides : int
    radius : float
    def _area(self): # concrete

class ArbitraryAperture(Aperture):
    binary_mask : int
    def _area(self):
        return self.binary_mask

class PhasePlate(Optic):
    def _phase_shift(self): # abstract
    def __call__(self): # concrete

class ArbitraryPhasePlate(Optic):
    phase : array
    def _phase_shift(self):
        return self.phase

class ChromaticPhasePlate(Optic):
    def _phase_shift(self, wavefront): # concrete
        # dependence on wavelength 
```
Although I was thinking about #60 and thought we could extend this even further by adding:
```py
class Source(Layer): # abstract
    def _create_wavefront(self): # abstract
    def __call__(self, parameters): # concrete

class CoherentPointSource(Layer):
    def _create_wavefront(self): # concrete

class CoherentPlaneSource(Layer): # Point source at infinty?
    def _create_wavefront(self): # concrete
```
where there are many other sources that could be generated. This leads to the natural extension of `CreateWavefront` into a set of many simple layers with specific properties. To me, it makes most sense to have `CoherentPointSource` output an `AngularWavefront` since the `Wavefront` is spherical and is best described by a field of view. `CoherentPlaneSource` on the other hand is better suited to a `PhysicalWavefront`. `GaussianWavefront` is a bit trickier because it requires a `Lens` or `Mirror` with a curvature. 

So my initial thoughts are included below. This is what @LouisDesdoigts and I originally discussed:
```py
class Layer(eqx.Module):
    def __call__(self, parameters): # abstract

class CurvedOptic(Layer):
    def _curvature(self): # abstract

class Mirror(Layer):
    def _reflect(self, wavefront): # abstract
    def __call__(self, parameters): # concrete

class GaussianMirror(CurvedOptic, Mirror):
    def _curvature(self): # concrete
    def _reflect(self, wavefront): # concrete

class HyperbolicMirror(CurvedOptic, Mirror):
    def _curvature(self): # concrete
    def _reflect(self, wavefront): # concrete

class SphericalMirror(CurvedOptic, Mirror):
    def _curvature(self): # concrete
    def _reflect(self, wavefront): # concrete

class Lens(CurvedOptic):
    index_of_refraction : float 
    def _curvature(self): # abstract
    def _focus(self, wavefront): # concrete
    def __call__(self, parameters): # concrete

class GaussianLens(Lens):
    def _curvature(self): # concrete

class HyperbolicLens(Lens):
    def _curvature(self): # concrete 

class ParabolicLens(Lens):
    def _curvature(self): # concrete

class Aperture(Optic):
    def _area(self): # Still abtract 
    def __call__(self): # concrete

class CircularAperture(Aperture):
    radius : float
    def _area(self, wavefront):

class RegularPolygonalAperture(Aperture):
    number_of_sides : int
    radius : float
    def _area(self): # concrete

class ArbitraryAperture(Aperture):
    binary_mask : int
    def _area(self):
        return self.binary_mask

class PhasePlate(Optic):
    def _phase_shift(self): # abstract
    def __call__(self): # concrete

class ArbitraryPhasePlate(Optic):
    phase : array
    def _phase_shift(self):
        return self.phase

class ChromaticPhasePlate(Optic):
    def _phase_shift(self, wavefront): # concrete
        # dependence on wavelength 
```
Although I was thinking about #60 and thought we could extend this even further by adding:
```py
class Source(Layer): # abstract
    def _create_wavefront(self): # abstract
    def __call__(self, parameters): # concrete

class CoherentPointSource(Layer):
    def _create_wavefront(self): # concrete

class CoherentPlaneSource(Layer): # Point source at infinty?
    def _create_wavefront(self): # concrete
```
where there are many other sources that could be generated. This leads to the natural extension of `CreateWavefront` into a set of many simple layers with specific properties. To me, it makes most sense to have `CoherentPointSource` output an `AngularWavefront` since the `Wavefront` is spherical and is best described by a field of view. `CoherentPlaneSource` on the other hand is better suited to a `PhysicalWavefront`. `GaussianWavefront` is a bit trickier because it requires a `Lens` or `Mirror` with a curvature. 



Fix patch:
diff --git a/.gitignore b/.gitignore
index f85a619d..423bb59a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,5 +1,8 @@
-src/__pycache__
 my_hours.txt
 site/
 tests/__pycache__
 https:/
+dLux/__pycache__
+dLux/utils/__pycache__
+dLux/dev/__pycache__
+dLux/dev/.ipynb_checkpoints
diff --git a/README.md b/README.md
index 38f961da..84815ad0 100644
--- a/README.md
+++ b/README.md
@@ -1,12 +1,7 @@
-
-∂Lux  Docs
-
----
-
 # ∂Lux
 [![PyPI version](https://badge.fury.io/py/dLux.svg)](https://badge.fury.io/py/dLux)
 [![License](https://img.shields.io/badge/License-BSD%203--Clause-blue.svg)](https://opensource.org/licenses/BSD-3-Clause)
-[![integration](https://github.com/LouisDesdoigts/dLux/actions/workflows/automated_tests.yml/badge.svg)](https://github.com/LouisDesdoigts/dLux/actions/workflows/automated_tests.yml)
+[![integration](https://github.com/LouisDesdoigts/dLux/actions/workflows/tests.yml/badge.svg)](https://github.com/LouisDesdoigts/dLux/actions/workflows/tests.yml)
 
 ∂Lux: Taking derivatives through Light - 'Optical systems as a Neural Network'
 
diff --git a/dLux/__init__.py b/dLux/__init__.py
new file mode 100644
index 00000000..871d6f3c
--- /dev/null
+++ b/dLux/__init__.py
@@ -0,0 +1,12 @@
+name = "dLux"
+
+from .base import OpticalSystem
+from .detectors import (ApplySaturation, ApplyPixelResponse, 
+    ApplyJitter)
+from .layers import (AddPhase, ApplyAperture, ApplyBasisCLIMB, 
+    ApplyBasisOPD, ApplyOPD, CircularAperture, CreateWavefront,
+    NormaliseWavefront, TiltWavefront)
+from .wavefronts import (GaussianWavefront, PhysicalWavefront, 
+    AngularWavefront)
+from .propagators import (PhysicalMFT, PhysicalFFT, AngularMFT, 
+    AngularFFT, PhysicalFresnel, GaussianPropagator)
diff --git a/src/base.py b/dLux/base.py
similarity index 97%
rename from src/base.py
rename to dLux/base.py
index 6934b15e..da71e5c6 100644
--- a/src/base.py
+++ b/dLux/base.py
@@ -3,7 +3,7 @@
 from jax import vmap
 import equinox as eqx
 from copy import deepcopy
-from dLux.wavefronts import *
+import dLux
 
 
 class OpticalSystem(eqx.Module):
@@ -91,7 +91,7 @@ def debug_prop(self, wavel, offset=np.zeros(2)):
         I believe this is diffable but there is no reason to force it to be
         """
         # params_dict = {"Wavefront": Wavefront(wavel, offset)}
-        params_dict = {"Wavefront": PhysicalWavefront(wavel, offset)}
+        params_dict = {"Wavefront": dLux.PhysicalWavefront(wavel, offset)}
         intermeds = []
         layers_applied = []
         for i in range(len(self.layers)):
@@ -145,7 +145,7 @@ def propagate_mono(self, wavel, offset=np.zeros(2)):
         
         """
         # params_dict = {"Wavefront": Wavefront(wavel, offset)}
-        params_dict = {"Wavefront": PhysicalWavefront(wavel, offset), 
+        params_dict = {"Wavefront": dLux.PhysicalWavefront(wavel, offset), 
                        "Optical System":self}
         
         for i in range(len(self.layers)):
diff --git a/src/detector_layers.py b/dLux/detectors.py
similarity index 100%
rename from src/detector_layers.py
rename to dLux/detectors.py
diff --git a/dLux/dev/layers.py b/dLux/dev/layers.py
new file mode 100644
index 00000000..3101ae1f
--- /dev/null
+++ b/dLux/dev/layers.py
@@ -0,0 +1,1041 @@
+"""
+src/dev/layers.py
+-----------------
+Development script for the new layers structure.
+"""
+
+# NOTE: Experimental code by @Jordan-Dennis is below.
+import equinox as eqx
+import jax.numpy as np
+import jax 
+import functools
+
+
+from typing import TypeVar, Dict
+from dLux.utils import (get_radial_positions, get_pixel_vector, 
+    get_pixel_positions)
+from abc import ABC, abstractmethod 
+
+
+Array = TypeVar("Array")
+Tensor = TypeVar("Tensor")
+Layer = TypeVar("Layer")
+Matrix = TypeVar("Matrix")
+
+
+MAX_DIFF = 4
+
+
+def cartesian_to_polar(coordinates : Tensor) -> Tensor:
+    """
+    Change the coordinate system from rectilinear to curvilinear.
+    
+    Parameters
+    ----------
+    coordinates : Tensor
+        The rectilinear coordinates.
+
+    Returns
+    -------
+    coordinates : Tensor
+        The curvilinear coordinates.
+    """
+    rho = np.hypot(coordinates[0], coordinates[1])
+    theta = np.arctan2(coordinates[1], coordinates[0])
+    return np.array([rho, theta])
+
+
+def factorial(n : int) -> int:
+    """
+    Calculate n! in a jax friendly way. Note that n == 0 is not a 
+    safe case.  
+
+    Parameters
+    ----------
+    n : int
+        The integer to calculate the factorial of.
+
+    Returns
+    n! : int
+        The factorial of the integer
+    """
+    return jax.lax.exp(jax.lax.lgamma(n + 1.))
+
+
+class Aperture(eqx.Module, ABC):
+    """
+    An abstract class that defines the structure of all the concrete
+    apertures. An aperture is represented by an array, usually in the
+    range of 0. to 1.. Values in between can be used to represent 
+    soft edged apertures and intermediate surfaces. 
+
+    Attributes
+    ----------
+    _npix : int
+        The number of pixels along one edge of the array which 
+        represents the aperture.
+    """
+    _npix : int
+    
+
+    def __init__(self : Layer, number_of_pixels : int) -> Layer:
+        """
+        Parameters
+        ----------
+        number_of_pixels : int
+            The number of pixels along one side of the array that 
+            represents this aperture.
+        """
+        self._npix = number_of_pixels
+
+    
+    @abstractmethod
+    def _aperture(self : Layer, number_of_pixels : int) -> Array:
+        """
+        Generate the aperture array as an array. 
+
+        Returns
+        -------
+        aperture : Array[Float]
+            The aperture. If these values are confined between 0. and 1.
+            then the physical interpretation is the transmission 
+            coefficient of that pixel. 
+        """
+
+
+    def get_npix(self : Layer) -> int:
+        """
+        Returns
+        -------
+        pixels : int
+            The number of pixels that parametrise this aperture.
+        """
+        return self._npix
+
+
+    def __call__(self : Layer, parameters : Dict) -> Dict:
+        """
+        Apply the aperture to an incoming wavefront.
+
+        Parameters
+        ----------
+        parameters : Dict
+            A dictionary containing the parameters of the model. 
+            The dictionary must satisfy `parameters.get("Wavefront")
+            != None`. 
+
+        Returns
+        -------
+        parameters : Dict
+            The parameter, parameters, with the "Wavefront"; key
+            value updated. 
+        """
+        wavefront = parameters["Wavefront"]
+        wavefront = wavefront.mulitply_amplitude(
+           self._aperture(self.get_npix()))
+        parameters["Wavefront"] = wavefront
+        return parameters
+
+
+class AnnularAperture(Aperture):
+    """
+    A circular aperture, parametrised by the number of pixels in
+    the array. By default this is a hard edged aperture but may be 
+    in future modifed to provide soft edges. 
+
+    Attributes
+    ----------
+    rmax : float
+        The proportion of the pixel vector that is contained within
+        the outer ring of the aperture.
+    rmin : float
+        The proportion of the pixel vector that is contained within
+        the inner ring of the aperture. 
+    """
+    rmin : float
+    rmax : float
+
+
+    def __init__(self : Layer, npix : int, rmax : float = 1., 
+            rmin : float = 0.) -> Layer:
+        """
+        Parameters
+        ----------
+        npix : int
+            The number of layers along one edge of the array that 
+            represents this aperture.
+        rmax : float = 1. 
+            The proportion of the pixel vector contained within the 
+            outer ring.
+        rmin : float = 0.
+            The proportion of the pixel vector contained within the
+            inner ring.
+        """
+        super().__init__(npix)
+        self.rmax = rmax
+        self.rmin = rmin
+
+
+    def _aperture(self : Layer) -> Array:
+        """
+        Generates an array representing a hard edged circular aperture.
+        All the values are 0. except for the outer edge. The t
+ 
+        Returns
+        -------
+        aperture : Array[Float]
+            The aperture. If these values are confined between 0. and 1.
+            then the physical interpretation is the transmission 
+            coefficient of that pixel. 
+        """
+        centre = (self.get_npix() - 1.) / 2.
+        coords = 2 / self.get_npix() * get_radial_positions(self.get_npix())
+        return np.logical_and(coords <= rmax, coords > rmin).astype(float)
+
+
+class HexagonalAperture(Aperture):
+    """
+    Generate a hexagonal aperture, parametrised by rmax. 
+    
+    Attributes
+    ----------
+    rmax : float
+        The proportion of the pixel vector that is taken up by a 
+        circle containing the hexagon.
+    x : int
+        The pixel coordinate in the x direction of the centre of the 
+        hexagonal aperture.
+    y : int
+        The pixel coordinate in the y direction of the centre of the 
+        hexagonal aperture. 
+    """
+    _rmax : float
+    _x : int
+    _y : int
+
+
+    def __init__(self : Layer, 
+            number_of_pixels_on_side : int, 
+            proportional_maximum_radius : float,
+            x_position_of_centre : int = 0, 
+            y_position_of_centre : int = 0) -> Layer:
+        """
+        Parameters
+        ----------
+        number_of_pixels_on_side : int
+            The number of pixels along one side of the square array
+        proportional_maximum_radius : float 
+            The outer radius of the smallest circle that contains the 
+            hexagon.
+        x_position_of_centre : int = 0
+            The x position of the centre of the hexagonal aperture.
+        y_poxition_of_centre : int = 0
+            The y position of the centre of the hexagonal aperture.             
+        """
+        super().__init__(number_of_pixels_on_side)
+        self._rmax = proportional_maximum_radius
+        self._x = x_position_of_centre
+        self._y = y_position_of_centre
+
+
+    def _get_centre(self : Layer) -> tuple:
+        """
+        The x and y coordinates of the centre of the hexagonal aperture.
+        
+        Returns
+        -------
+        x, y : tuple
+            The x and y coordinates of the centre in pixels
+        """
+        return self._x, self._y
+
+
+    def _get_rmax(self : Layer) -> float:
+        """
+        Returns
+        -------
+        rmax : float
+            The proportional maximum radius of the smallest circle that
+            can completely enclose the aperture.
+        """
+        return self._rmax
+
+
+    def _aperture(self : Layer) -> Array:
+        """
+        Generates an array representing the hard edged hexagonal 
+        aperture. 
+
+        Returns
+        -------
+        aperture : Array
+            The aperture represented as a binary float array of 0. and
+            1. representing no transmission and transmission 
+            respectively.
+        """
+        x_centre, y_centre = self._get_centre()
+        number_of_pixels = self._get_npix()
+        maximum_radius = self._get_rmax()
+
+        x, y = _get_pixel_positions(number_of_pixels, -x_centre,
+            -y_centre)
+
+        x *= 2 / number_of_pixels
+        y *= 2 / number_of_pixels
+
+        rectangle = (np.abs(x) <= maximum_radius / 2.) \
+            & (np.abs(y) <= (maximum_radius * np.sqrt(3) / 2.))
+
+        left_triangle = (x <= - maximum_radius / 2.) \
+            & (x >= - maximum_radius) \
+            & (np.abs(y) <= (x + maximum_radius) * np.sqrt(3))
+
+        right_triangle = (x >= maximum_radius / 2.) \
+            & (x <= maximum_radius) \
+            & (np.abs(y) <= (maximum_radius - x) * np.sqrt(3))
+
+        hexagon = rectangle | left_triangle | right_triangle
+        return np.asarray(hexagon).astype(float)
+
+
+class BasisPhase(eqx.Module):
+    """
+    _Abstract_ class representing a basis fixed over an aperture 
+    that is used to optimise and learn aberations in the aperture. 
+    
+    Attributes
+    ----------
+    _x : float
+        The x coordinate of the centre of the aperture.
+    _y : float
+        The y coordinate of the centre of the aperture.
+    _r : float
+        The radius of the aperture. The radius belongs to the smallest
+        circle that completely contains the aperture. For the math
+        nerds the infimum of the set of circles whose union with the
+        aperture is the aperture. 
+    _npix : int
+        The number of pixels along the edge of the square array 
+        representing each term in the basis.
+    """
+    x : float
+    y : float
+    rmax : float
+    npix : int
+    nterms : int
+
+
+    def __init__(self : Layer, x : float, y : float, r : float,
+            npix : int, nterms : int) -> Layer:
+        """
+        Parameters
+        ----------
+        x : float
+            The x coordinate of the centre of the aperture containing 
+            the basis.
+        y : float
+            The y coordinate of the centre of the aperture containing 
+            the basis.
+        r : float
+            The radius of the aperture containing the basis. For 
+            abitrary shapes the radius is belongs to the smallest 
+            circle that can completely contain the aperture. 
+        npix : int
+            The number of pixels along one side of the basis arrays.
+            That is each term in the basis will be evaluated on a 
+            `npix` by `npix` grid.
+        nterms : int
+            The number of basis terms to generate. This determines the
+            length of the leading dimension of the output Tensor. 
+        """
+        self.x = np.asarray(x).astype(float)
+        self.y = np.asarray(y).astype(float)
+        self.rmax = np.asarray(r).astype(float)
+        self.npix = int(npix)
+        self.nterms = int(nterms)
+
+
+    @abstractmethod
+    def _aperture(self : Layer) -> Matrix:
+        """
+        Generate a binary or soft edged aperture mask in which to 
+        construct the basis polynomials. 
+
+        Returns
+        -------
+        aperture : Matrix
+            The aperture encoded in an `npix` by `npix` array. Note 
+            that with hard edged apertures the gradient can become 
+            unusual.
+        """
+
+
+    @abstractmethod
+    @functools.cached_property 
+    def _basis(self : Layer, n : int) -> Tensor:
+        """
+        Generate the basis. Requires a single run after which,
+        the basis is cached and can be used with no computational 
+        cost.  
+
+        Parameters
+        ----------
+        n : int 
+            The number of terms in the basis to generate. 
+
+        Returns
+        -------
+        basis : Tensor
+            The basis polynomials evaluated on the square arrays
+            containing the apertures until `maximum_radius`.
+            The leading dimension is `n` long and 
+            each stacked array is a basis term. The final shape is:
+            ```py
+            basis.shape == (n, npix, npix)
+            ```
+        """
+
+
+    def save(self : Layer, file_name : str, n : int) -> None:
+        """
+        Save the basis to a file.
+
+        Parameters
+        ----------
+        file_name : str
+            The name of the file to save the basis terms to.
+        n : int
+            The number of terms in the basis to generate in the save.
+        """
+        basis = self._basis
+        with open(file_name, "w") as save:
+            save.write(basis)
+
+
+    def __call__(self : Layer, parameters : dict) -> dict:
+        """
+        Apply a phase shift to the wavefront based on the basis 
+        terms that have been generated.
+
+        Overrides
+        ---------
+        __call__ : Layer
+            Provides a concrete implementation of the `__call__` method
+            defined in the abstract base class `Layer`.
+
+        Parameters
+        ----------
+        parameters : Dict
+            A dictionary containing the parameters of the model. 
+            The dictionary must satisfy `parameters.get("Wavefront")
+            != None`. 
+
+        Returns
+        -------
+        parameters : Dict
+            The parameter, parameters, with the "Wavefront"; key
+            value updated. 
+        """
+        wavefront = parameters["Wavefront"]
+        parameters["Wavefront"] = wavefront.add_phase(self._basis)
+        return parameters
+        
+
+class PolygonalBasis(BasisPhase, ABC):
+    """
+    Orthonormalises the zernike basis over a polygon to provide
+    a basis on a polygonal aperture.
+
+    Attributes
+    ----------
+    theta : float, radians
+        This represents the angle of rotation from the positive x 
+        axis. 
+    phi : float, radians
+        This represents the angle of shear measured from the positive 
+        y axis.
+    """
+    theta : float
+    phi : float
+
+
+    def __init__(self : Layer, nterms : int, npix : int, 
+            rmax : float, theta : float, phi : float, x : float,
+            y : float) -> Layer:
+        """
+        Parameters
+        ---------- 
+        nterm : int 
+            The number of polynomials to generate. This is not a 
+            gradable parameter.
+        npix : int
+            The number of pixels in the image that is to be generated.
+        rmax : float
+            The radius of the smallest circle that can fully enclose the 
+            aperture upon which the basis will be orthonormalised. The 
+            coordinates hadnled internally are normalised, but this 
+            quantity can be any number. > 1 is a magnification and < 1
+            makes it smaller. 
+        theta : float, radians
+            This represents the angle of rotation from the positive x 
+            axis. 
+        phi : float, radians
+            This represents the angle of shear measured from the positive 
+            y axis.
+        x : float
+            The _x_ coordinates of the centre of the aperture. This 
+            occurs following the normalisation so that the type can 
+            be given as a `float` for gradient stability.
+        y : float
+            The _y_ coordinates of the centre of the aperture. This 
+            occurs following the normalisation so that the type can 
+            be given as a `float` for gradient stability.
+        """
+        super().__init__(x, y, rmax, npix, nterms)
+        self.theta = np.asarray(theta).astype(float)
+        self.phi = np.asarray(phi).astype(float)
+
+
+    @functools.partial(jax.vmap, in_axes=(None, 0))
+    def _noll_index(self : Layer, j : int) -> tuple:
+        """
+        Decode the jth noll index of the zernike polynomials. This 
+        arrises because the zernike polynomials are parametrised by 
+        a pair numbers, e.g. n, m, but we want to impose an order.
+        The noll indices are the standard way to do this see [this]
+        (https://oeis.org/A176988) for more detail. The top of the 
+        mapping between the noll index and the pair of numbers is 
+        shown below:
+
+        n, m Indices
+        ------------
+        (0, 0)
+        (1, -1), (1, 1)
+        (2, -2), (2, 0), (2, 2)
+        (3, -3), (3, -1), (3, 1), (3, 3)
+
+        Noll Indices
+        ------------
+        1
+        3, 2
+        5, 4, 6
+        9, 7, 8, 10
+
+        Parameters
+        ----------
+        j : int
+            The noll index to decode.
+        
+        Returns
+        -------
+        n, m : tuple
+            The n, m parameters of the zernike polynomial.
+        """
+        # To retrive the row that we are in we use the formula for 
+        # the sum of the integers:
+        #  
+        #  n      n(n + 1)
+        # sum i = -------- = x_{n}
+        # i=0        2
+        # 
+        # However, `j` is a number between x_{n - 1} and x_{n} to 
+        # retrieve the 0th based index we want the upper bound. 
+        # Applying the quadratic formula:
+        # 
+        # n = -1/2 + sqrt(1 + 8x_{n})/2
+        #
+        # We know that n is an integer and hence of x_{n} -> j where 
+        # j is not an exact solution the row can be found by taking 
+        # the floor of the calculation. 
+        #
+        # n = (-1/2 + sqrt(1 + 8j)/2) // 1
+        #
+        # All the odd noll indices map to negative m integers and also 
+        # 0. The sign can therefore be determined by -(j & 1). 
+        # This works because (j & 1) returns the rightmost bit in 
+        # binary representation of j. This is equivalent to -(j % 2).
+        # 
+        # The m indices range from -n to n in increments of 2. The last 
+        # thing to do is work out how many times to add two to -n. 
+        # This can be done by banding j away from the smallest j in 
+        # the row. 
+        #
+        # The smallest j in the row can be calculated using the sum of
+        # integers formula in the comments above with n = (n - 1) and
+        # then adding one. Let this number be (x_{n - 1} + 1). We can 
+        # then subtract j from it to get r = (j - x_{n - 1} + 1)
+        #
+        # The odd and even cases work differently. I have included the 
+        # formula below:
+        # odd : p = (j - x_{n - 1}) // 2 
+       
+        # even: p = (j - x_{n - 1} + 1) // 2
+        # where p represents the number of times 2 needs to be added
+        # to the base case. The 1 required for the even case can be 
+        # generated in place using ~(j & 1) + 2, which is 1 for all 
+        # even numbers and 0 for all odd numbers.
+        #
+        # For odd n the base case is 1 and for even n it is 0. This 
+        # is the result of the bitwise operation j & 1 or alternatively
+        # (j % 2). The final thing is adding the sign to m which is 
+        # determined by whether j is even or odd hence -(j & 1).
+        n = (np.ceil(-1 / 2 + np.sqrt(1 + 8 * j) / 2) - 1).astype(int)
+        smallest_j_in_row = n * (n + 1) / 2 + 1 
+        number_of_shifts = (j - smallest_j_in_row + ~(n & 1) + 2) // 2
+        sign_of_shift = -(j & 1) + ~(j & 1) + 2
+        base_case = (n & 1)
+        m = (sign_of_shift * (base_case + number_of_shifts * 2)).astype(int)
+        return n, m
+
+
+    def _radial_zernike(self : Layer, n : int, m : int,
+            rho : Matrix) -> Tensor:
+        """
+        The radial zernike polynomial.
+
+        Parameters
+        ----------
+        n : int
+            The first index number of the zernike polynomial to forge
+        m : int 
+            The second index number of the zernike polynomial to forge.
+        rho : Matrix
+            The radial positions of the aperture. Passed as an argument 
+            for speed.
+
+        Returns
+        -------
+        radial : Tensor
+            An npix by npix stack of radial zernike polynomials.
+        """
+        m, n = np.abs(m), np.abs(n)
+        upper = ((np.abs(n) - np.abs(m)) / 2).astype(int) + 1
+        rho = np.tile(rho, (MAX_DIFF, 1, 1))
+
+        murder_weapon = (np.arange(MAX_DIFF) < upper)
+
+        k = np.arange(MAX_DIFF) * murder_weapon
+        coefficients = (-1) ** k * factorial(n - k) / \
+            (factorial(k) * \
+                factorial(((n + m) / 2).astype(int) - k) * \
+                factorial(((n - m) / 2).astype(int) - k))
+        radial = coefficients.reshape(MAX_DIFF, 1, 1) *\
+            rho ** (n - 2 * k).reshape(MAX_DIFF, 1, 1) *\
+            murder_weapon.reshape(MAX_DIFF, 1, 1)
+         
+        return radial.sum(axis=0)
+
+
+    def _zernikes(self : Layer, coordinates : Tensor) -> Tensor:
+        """
+        Calculate the zernike basis on a square pixel grid. 
+
+        Parameters
+        ----------
+        number : int
+            The number of zernike basis terms to calculate.
+            This is a static argument to jit because the array
+            size depends on it.
+        pixels : int
+            The number of pixels along one side of the zernike image
+            for each of the n zernike polynomials.
+        coordinates : Tensor
+            The cartesian coordinates to generate the hexikes on.
+            The dimensions of the tensor should be `(2, npix, npix)`.
+            where the leading axis is the x and y dimensions.  
+
+        Returns
+        -------
+        zernike : Tensor 
+            The zernike polynomials evaluated until number. The shape
+            of the output tensor is number by pixels by pixels. 
+        """
+        j = np.arange(1, self.nterms + 1).astype(int)
+        n, m = self._noll_index(j)
+        coordinates = cartesian_to_polar(coordinates)
+
+        # NOTE: The idea is to generate them here at the higher level 
+        # where things will not change and we will be done. 
+        rho = coordinates[0]
+        theta = coordinates[1]
+
+        aperture = (rho <= 1.).astype(int)
+
+        # In the calculation of the noll coefficient we must define 
+        # between the m == 0 and and the m != 0 case. I have done 
+        # this in place by casting the logical operation to an int. 
+
+        normalisation_coefficients = \
+            (1 + (np.sqrt(2) - 1) * (m != 0).astype(int)) \
+            * np.sqrt(n + 1)
+
+        radial_zernikes = np.zeros((self.nterms,) + rho.shape)
+        for i in np.arange(self.nterms):
+            radial_zernikes = radial_zernikes\
+                .at[i]\
+                .set(self._radial_zernike(n[i], m[i], rho))
+
+        # When m < 0 we have the odd zernike polynomials which are 
+        # the radial zernike polynomials multiplied by a sine term.
+        # When m > 0 we have the even sernike polynomials which are 
+        # the radial polynomials multiplies by a cosine term. 
+        # To produce this result without logic we can use the fact
+        # that sine and cosine are separated by a phase of pi / 2
+        # hence by casting int(m < 0) we can add the nessecary phase.
+        out_shape = (self.nterms, 1, 1)
+
+        theta = np.tile(theta, out_shape)
+        m = m.reshape(out_shape)
+        phase_mod = (m < 0).astype(int) * np.pi / 2
+        phase = np.cos(np.abs(m) * theta - phase_mod)
+
+        normalisation_coefficients = \
+            normalisation_coefficients.reshape(out_shape)
+        
+        return normalisation_coefficients * radial_zernikes \
+            * aperture * phase 
+
+
+    def _orthonormalise(self : Layer, aperture : Matrix, 
+            zernikes : Tensor) -> Tensor:
+        """
+        The hexike polynomials up until `number_of_hexikes` on a square
+        array that `number_of_pixels` by `number_of_pixels`. The 
+        polynomials can be restricted to a smaller subset of the 
+        array by passing an explicit `maximum_radius`. The polynomial
+        will then be defined on the largest hexagon that fits with a 
+        circle of radius `maximum_radius`. 
+        
+        Parameters
+        ----------
+        aperture : Matrix
+            An array representing the aperture. This should be an 
+            `(npix, npix)` array. 
+        number_of_hexikes : int = 15
+            The number of basis terms to generate. 
+        zernikes : Tensor
+            The zernike polynomials to orthonormalise on the aperture.
+            This tensor should be `(nterms, npix, npix)` in size, where 
+            the first axis represents the noll indexes. 
+
+        Returns
+        -------
+        hexikes : Tensor
+            The hexike polynomials evaluated on the square arrays
+            containing the hexagonal apertures until `maximum_radius`.
+            The leading dimension is `number_of_hexikes` long and 
+            each stacked array is a basis term. The final shape is:
+            ```py
+            hexikes.shape == (number_of_hexikes, number_of_pixels, number_of_pixels)
+            ```
+        """
+        pixel_area = aperture.sum()
+        basis = np.zeros(zernikes.shape).at[0].set(aperture)
+        
+        for j in np.arange(1, self.nterms):
+            intermediate = zernikes[j] * aperture
+
+            coefficient = -1 / pixel_area * \
+                (zernikes[j] * basis[1 : j + 1] * aperture)\
+                .sum(axis = (1, 2))\
+                .reshape(j, 1, 1) 
+
+            intermediate += (coefficient * basis[1 : j + 1])\
+                .sum(axis = 0)
+            
+            basis = basis\
+                .at[j]\
+                .set(intermediate / \
+                    np.sqrt((intermediate ** 2).sum() / pixel_area))
+        
+        return basis
+
+
+    def _magnify(self : Layer, coordinates : Tensor) -> Tensor:
+        """
+        Enlarge or shrink the coordinate system, by the inbuilt 
+        amount specified by `self._rmax`.
+
+        Parameters
+        ----------
+        coordinates : Tensor
+            A `(2, npix, npix)` representation of the coordinate 
+            system. The leading dimensions specifies the x and then 
+            the y coordinates in that order. 
+
+        Returns
+        -------
+        coordinates : Tensor
+            The enlarged or shrunken coordinate system.
+        """
+        return 1 / self.rmax * coordinates
+
+
+    def _rotate(self : Layer, coordinates : Tensor) -> Tensor:
+        """
+        Rotate the coordinate system by a pre-specified amount,
+        `self._theta`
+
+        Parameters
+        ----------
+        coordinates : Tensor
+            A `(2, npix, npix)` representation of the coordinate 
+            system. The leading dimensions specifies the x and then 
+            the y coordinates in that order. 
+
+        Returns
+        -------
+        coordinates : Tensor
+            The rotated coordinate system. 
+        """
+        rotation_matrix = np.array([
+            [np.cos(self.theta), -np.sin(self.theta)],
+            [np.sin(self.theta), np.cos(self.theta)]])            
+        return np.apply_along_axis(np.matmul, 0, coordinates, 
+            rotation_matrix) 
+
+
+    def _shear(self : Layer, coordinates : Tensor) -> Tensor:
+        """
+        Shear the coordinate system by the inbuilt amount `self._phi`.
+
+        Parameters
+        ----------
+        coordinates : Tensor
+            A `(2, npix, npix)` representation of the coordinate 
+            system. The leading dimensions specifies the x and then 
+            the y coordinates in that order. 
+
+        Returns
+        -------
+        coordinates : Tensor
+            The sheared coordinate system. 
+        """
+        return coordinates\
+            .at[0]\
+            .set(coordinates[0] - coordinates[1] * np.tan(self.phi)) 
+
+
+    def _offset(self : Layer, coordinates : Tensor) -> Tensor:
+        """
+        Offset the coordinate system by prespecified amounts in both
+        the `x` and `y` directions. 
+
+        Parameters
+        ----------
+        coordinates : Tensor
+            A `(2, npix, npix)` representation of the coordinate 
+            system. The leading dimensions specifies the x and then 
+            the y coordinates in that order. 
+
+        Returns
+        -------
+        coordinates : Tensor
+            The translated coordinate system. 
+        """
+        return coordinates\
+            .at[0]\
+            .set(coordinates[0] - self.x)\
+            .at[1]\
+            .set(coordinates[1] - self.y)
+
+
+    def _basis(self : Layer):
+        """
+        Generate the basis. Requires a single run after which,
+        the basis is cached and can be used with no computational 
+        cost.  
+
+        Returns
+        -------
+        basis : Tensor
+            The basis polynomials evaluated on the square arrays
+            containing the apertures until `maximum_radius`.
+            The leading dimension is `n` long and 
+            each stacked array is a basis term. The final shape is:
+            `(n, npix, npix)`
+        """
+        coordinates = self._coordinates()
+        zernikes = self._zernikes(coordinates)
+        aperture = self._aperture(coordinates)
+
+        return self._orthonormalise(aperture, zernikes)  
+
+
+    def _coordinates(self : Layer) -> Tensor:
+        """
+        Generate the transformed coordinate system for the aperture.
+
+        Returns
+        -------
+        coordinates : Tensor
+            The coordinate system in the rectilinear view, with the
+            x and y coordinates stacked above one another.
+        """
+        coordinates = self._shear(
+            self._rotate(
+                self._offset(
+                    self._magnify(
+                        2 / self.npix * get_pixel_positions(self.npix)))))
+        return coordinates
+
+
+    def set_theta(self : Layer, theta : float) -> Layer:
+        """
+        Parameters
+        ----------
+        theta : float
+            The angle of rotation from the positive x-axis.  
+
+        Returns
+        -------
+        basis : HexagonalBasis 
+            The rotated hexagonal basis. 
+        """
+        return eqx.tree_at(lambda basis : basis.theta, self, theta)
+
+
+    def set_magnification(self : Layer, rmax : float) -> Layer:
+        """
+        Parameters
+        ----------
+        rmax : float
+            The radius of the smallest circle that can completely 
+            enclose the aperture.
+
+        Returns
+        -------
+        basis : HexagonalBasis
+            The magnified hexagonal basis.
+        """
+        return eqx.tree_at(lambda basis : basis.rmax, self, rmax)
+
+
+    def set_shear(self : Layer, phi : float) -> Layer:
+        """
+        Parameters
+        ----------
+        phi : float
+            The angle of shear from the positive y-axis.
+
+        Returns
+        -------
+        basis : HexagonalBasis
+            The sheared hexagonal basis.
+        """
+        return eqx.tree_at(lambda basis : basis.phi, self, phi)      
+
+
+    def set_x_offset(self : Layer, x : float) -> Layer:
+        """
+        Parameters
+        ----------
+        x : float
+            The x coordinate of the centre of the hexagonal
+            aperture.
+
+        Returns
+        -------
+        basis : HexagonalBasis
+            The translated hexagonal basis. 
+        """
+        return eqx.tree_at(lambda basis : basis.x, self, x)
+
+
+    def set_y_offset(self : Layer, y : float) -> Layer:
+        """
+        Parameters
+        ----------
+        x : float
+            The y coordinate of the centre of the hexagonal
+            aperture.
+
+        Returns
+        -------
+        basis : HexagonalBasis
+            The translated hexagonal basis. 
+        """
+        return eqx.tree_at(lambda basis : basis.y, self, y)
+
+
+class HexagonalBasis(PolygonalBasis):
+    """
+    A basis on hexagonal surfaces. This is based on [this]
+    (https://github.com/spacetelescope/poppy/poppy/zernike.py)
+    code from the _POPPY_ library, but ported into [_JAX_]
+    (https://github.com/google/jax) and allowed to perform 
+    more complicated spatial transformations. 
+    """
+    def __init__(self : Layer, nterms : int, npix : int,
+            rmax : float, theta : float,
+            phi : float, x : float, y : float) -> Layer:
+        """
+        Parameters
+        ----------
+        nterms : int
+            The number of terms to generate in the basis. 
+        npix : int
+            The number of pixels in the image that is to be generated.
+        rmax : float
+            The radius of the smallest circle that can fully enclose the 
+            aperture upon which the basis will be orthonormalised. The 
+            coordinates hadnled internally are normalised, but this 
+            quantity can be any number. > 1 is a magnification and < 1
+            makes it smaller. 
+        theta : float, radians
+            This represents the angle of rotation from the positive x 
+            axis. 
+        phi : float, radians
+            This represents the angle of shear measured from the positive 
+            y axis.
+        x : float
+            The _x_ coordinates of the centre of the aperture. This 
+            occurs following the normalisation so that the type can 
+            be given as a `float` for gradient stability.
+        y : float
+            The _y_ coordinates of the centre of the aperture. This 
+            occurs following the normalisation so that the type can 
+            be given as a `float` for gradient stability.
+        """
+        super().__init__(nterms, npix, rmax, theta, phi, x, y)
+
+
+    def _aperture(self : Layer,
+            coordinates : Tensor,
+            maximum_radius : float = 1.) -> Array:
+        """
+        Generate a binary mask representing the pixels occupied by 
+        the aperture. 
+
+        Parameters
+        ----------
+        coordinates : Tensor
+            The coordinate grid as a stacked array that is 
+            `(2, npix, npix)`, where the leading axis denotes the x and 
+            y coordinate sets. This Tensor must be normalised, if using 
+            the inbuilt `get_pixel_positions` normalisation can be done
+            via `2 / npix`.
+        maximum_radius : float 
+            The radius of the smallest circle that can completely contain
+            the entire aperture. 
+
+        Returns
+        -------
+        aperture : Array
+            The bitmask that represents the circular aperture.        
+        """
+        x, y = coordinates[0], coordinates[1]
+
+        rectangle = (np.abs(x) <= maximum_radius / 2.) \
+            & (np.abs(y) <= (maximum_radius * np.sqrt(3) / 2.))
+
+        left_triangle = (x <= - maximum_radius / 2.) \
+            & (x >= - maximum_radius) \
+            & (np.abs(y) <= (x + maximum_radius) * np.sqrt(3))
+
+        right_triangle = (x >= maximum_radius / 2.) \
+            & (x <= maximum_radius) \
+            & (np.abs(y) <= (maximum_radius - x) * np.sqrt(3))
+
+        hexagon = rectangle | left_triangle | right_triangle
+        return np.asarray(hexagon).astype(float)
diff --git a/src/layers.py b/dLux/layers.py
similarity index 99%
rename from src/layers.py
rename to dLux/layers.py
index c889f982..de4995c1 100644
--- a/src/layers.py
+++ b/dLux/layers.py
@@ -700,3 +700,4 @@ def CLIMB(self, wf, ppsz = 256):
         soft_bin = vmap_mask(flat).reshape(ppsz, ppsz)
 
         return soft_bin
+
diff --git a/src/propagators.py b/dLux/propagators.py
similarity index 100%
rename from src/propagators.py
rename to dLux/propagators.py
diff --git a/dLux/utils/__init__.py b/dLux/utils/__init__.py
new file mode 100644
index 00000000..5f831e3c
--- /dev/null
+++ b/dLux/utils/__init__.py
@@ -0,0 +1,12 @@
+name = "utils"
+
+from .plotting import plot_batch
+from .zernike import zernike_basis
+from .hexike import hexike_basis
+from .helpers import (unif_rand, norm_rand, opd2phase, phase2opd,
+    get_ppf, rad2arcsec, arcsec2rad, cart2polar, polar2cart,
+    scale_mask, get_GE, get_RGE, get_RWGE, get_Rvec, get_Rmask,
+    nyquist_pix_airy, nyquist_pix)
+from .bayes import (poiss_logpri, chi2_logpri, calc_cov, calc_ent)
+from .coordinates import (get_pixel_vector, get_pixel_positions,
+    get_radial_positions)
diff --git a/dLux/utils/__pycache__/__init__.cpython-310.pyc b/dLux/utils/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 00000000..1fae7d0f
Binary files /dev/null and b/dLux/utils/__pycache__/__init__.cpython-310.pyc differ
diff --git a/dLux/utils/__pycache__/bayes.cpython-310.pyc b/dLux/utils/__pycache__/bayes.cpython-310.pyc
new file mode 100644
index 00000000..f7da3d92
Binary files /dev/null and b/dLux/utils/__pycache__/bayes.cpython-310.pyc differ
diff --git a/dLux/utils/__pycache__/coordinates.cpython-310.pyc b/dLux/utils/__pycache__/coordinates.cpython-310.pyc
new file mode 100644
index 00000000..e04168cf
Binary files /dev/null and b/dLux/utils/__pycache__/coordinates.cpython-310.pyc differ
diff --git a/dLux/utils/__pycache__/helpers.cpython-310.pyc b/dLux/utils/__pycache__/helpers.cpython-310.pyc
new file mode 100644
index 00000000..60e42f51
Binary files /dev/null and b/dLux/utils/__pycache__/helpers.cpython-310.pyc differ
diff --git a/dLux/utils/__pycache__/hexike.cpython-310.pyc b/dLux/utils/__pycache__/hexike.cpython-310.pyc
new file mode 100644
index 00000000..e0904d8a
Binary files /dev/null and b/dLux/utils/__pycache__/hexike.cpython-310.pyc differ
diff --git a/dLux/utils/__pycache__/plotting.cpython-310.pyc b/dLux/utils/__pycache__/plotting.cpython-310.pyc
new file mode 100644
index 00000000..481c7d6a
Binary files /dev/null and b/dLux/utils/__pycache__/plotting.cpython-310.pyc differ
diff --git a/dLux/utils/__pycache__/zernike.cpython-310.pyc b/dLux/utils/__pycache__/zernike.cpython-310.pyc
new file mode 100644
index 00000000..dcf281f7
Binary files /dev/null and b/dLux/utils/__pycache__/zernike.cpython-310.pyc differ
diff --git a/src/bayes.py b/dLux/utils/bayes.py
similarity index 100%
rename from src/bayes.py
rename to dLux/utils/bayes.py
diff --git a/dLux/utils/coordinates.py b/dLux/utils/coordinates.py
new file mode 100644
index 00000000..2abf2d05
--- /dev/null
+++ b/dLux/utils/coordinates.py
@@ -0,0 +1,102 @@
+"""
+src/utilities/image.py
+----------------------
+A static class that implements centering conventions for dealing 
+with para-axial arrays. 
+"""
+__author__ = "Jordan Dennis"
+__date__ = "07/07/2022"
+
+
+import jax.numpy as np
+from typing import TypeVar
+
+
+Vector = TypeVar("Vector")
+Matrix = TypeVar("Matrix")
+Tensor = TypeVar("Tensor")
+Array = TypeVar("Array")
+
+
+def get_pixel_vector(
+        number_of_pixels : int,
+        pixel_offset : int = 0) -> Array:
+    """
+    Generate the coordinates along the edge of the image.
+
+    Parameters
+    ----------
+    number_of_pixels : int
+        The number of pixels along the edge of the pixel plane.
+    pixel_offset : int
+        The number of pixels to offset the zero.
+
+    Returns
+    -------
+    x : Array
+        The pixel positions along the edge of the image.
+    """
+    return np.arange(number_of_pixels) - number_of_pixels / 2. \
+        + 0.5 + pixel_offset
+    
+
+def get_pixel_positions(
+        number_of_pixels : int, 
+        x_pixel_offset : int = 0,
+        y_pixel_offset : int = 0) -> Array:
+    """
+    Generates offset para-axial coordinates, defining the optical 
+    axis. 
+
+    Parameters
+    ----------
+    number_of_pixels : int
+        The number of pixels along the side of the output array.
+    x_pixel_offset : int = 0.
+        The x offset of the centre of the coordinate system in the 
+        square output array.
+    y_pixel_offset : int = 0
+        The y offset of the centre of the coordinate system in the 
+        square output array.
+
+    Returns 
+    -------
+    pixel_positions : Array
+        The pixel positions in the square output array with the 
+        correct offsets
+    """
+    x = get_pixel_vector(number_of_pixels, x_pixel_offset)
+    y = get_pixel_vector(number_of_pixels, y_pixel_offset)
+    return np.array(np.meshgrid(x, y))
+
+
+# TODO:     
+def get_radial_positions(number_of_pixels : int,
+        x_pixel_offset : int,
+        y_pixel_offset : int) -> Matrix:
+    """
+    Generate the radial coordinates of each pixel. 
+
+    Parameters
+    ----------
+    number_of_pixels : int
+        The number of pixels along one edge of the square array.
+    x_pixel_offset : int = 0.
+        The x offset of the centre of the coordinate system in the 
+        square output array.
+    y_pixel_offset : int = 0
+        The y offset of the centre of the coordinate system in the 
+        square output array.
+
+    Returns
+    -------
+    positions : Matrix
+        A grid of the radial coordinates.
+    """
+    # NOTE: I think CircularAperture is Broken 
+    positions = get_pixel_positions(number_of_pixels, 
+        x_pixel_offset, y_pixel_offset)
+    x, y = positions[0], positions[1]
+    rho = np.hypot(x, y)
+    theta = np.arctan2(y, x)
+    return np.array([rho, theta])
diff --git a/src/helpers.py b/dLux/utils/helpers.py
similarity index 100%
rename from src/helpers.py
rename to dLux/utils/helpers.py
diff --git a/dLux/utils/hexike.py b/dLux/utils/hexike.py
new file mode 100644
index 00000000..82c495a0
--- /dev/null
+++ b/dLux/utils/hexike.py
@@ -0,0 +1,213 @@
+"""
+An attempt at generating the hexike polynomials for optimising a 
+hexagonal surface for aberations. 
+"""
+
+import jax.numpy as np
+import dLux
+from typing import TypeVar
+
+Array = TypeVar("Array")
+Tensor = TypeVar("Tensor")
+
+
+def _get_pixel_vector(
+        number_of_pixels : int,
+        pixel_offset : int = 0) -> Array:
+    """
+    Generate the coordinates along the edge of the image.
+
+    Parameters
+    ----------
+    number_of_pixels : int
+        The number of pixels along the edge of the pixel plane.
+    pixel_offset : int
+        The number of pixels to offset the zero.
+
+    Returns
+    -------
+    x : Array
+        The pixel positions along the edge of the image.
+    """
+    return np.arange(number_of_pixels) - number_of_pixels / 2. \
+        + 0.5 + pixel_offset
+    
+
+def _get_pixel_positions(
+        number_of_pixels : int, 
+        x_pixel_offset : int = 0,
+        y_pixel_offset : int = 0) -> Array:
+    """
+    Generates offset para-axial coordinates, defining the optical 
+    axis. 
+
+    Parameters
+    ----------
+    number_of_pixels : int
+        The number of pixels along the side of the output array.
+    x_pixel_offset : int = 0.
+        The x offset of the centre of the coordinate system in the 
+        square output array.
+    y_pixel_offset : int = 0
+        The y offset of the centre of the coordinate system in the 
+        square output array.
+
+    Returns 
+    -------
+    pixel_positions : Array
+        The pixel positions in the square output array with the 
+        correct offsets
+    """
+    x = _get_pixel_vector(number_of_pixels, x_pixel_offset)
+    y = _get_pixel_vector(number_of_pixels, y_pixel_offset)
+    return np.meshgrid(x, y)
+
+
+
+# TODO: Work out how to stop this repetition of logic. 
+def _hexagonal_aperture(
+        number_of_pixels : int, 
+        x_pixel_offset : int,
+        y_pixel_offset : int,
+        maximum_radius : float = 1.) -> Array:
+    """
+    Generate a binary mask representing the pixels occupied by 
+    the aperture. 
+
+    Parameters
+    ----------
+    number_of_pixels : int
+        The ouput array will have the shape `aperture.shape == 
+        (number_of_pixels, number_of_pixels)`. 
+    x_pixel_offset : int
+        The offset of the aperture in the square output array in the 
+        x direction.
+    y_pixel_offset : int
+        The offset of the aperture in the square output array in the 
+        y direction. 
+    maximum_radius : float 
+        The radius of the smallest circle that can completely contain
+        the entire aperture. 
+
+    Returns
+    -------
+    aperture : Array
+        The bitmask that represents the circular aperture.        
+    """
+    # NOTE: This is where I need to choose generate everything correctly 
+    # I think. 
+    x, y = _get_pixel_positions(number_of_pixels, x_pixel_offset,
+        y_pixel_offset)
+
+    x *= 2 / number_of_pixels
+    y *= 2 / number_of_pixels
+
+    rectangle = (np.abs(x) <= maximum_radius / 2.) \
+        & (np.abs(y) <= (maximum_radius * np.sqrt(3) / 2.))
+
+    left_triangle = (x <= - maximum_radius / 2.) \
+        & (x >= - maximum_radius) \
+        & (np.abs(y) <= (x + maximum_radius) * np.sqrt(3))
+
+    right_triangle = (x >= maximum_radius / 2.) \
+        & (x <= maximum_radius) \
+        & (np.abs(y) <= (maximum_radius - x) * np.sqrt(3))
+
+    hexagon = rectangle | left_triangle | right_triangle
+    return np.asarray(hexagon).astype(float)
+
+
+def hexike_basis(
+        number_of_hexikes : int = 15, 
+        number_of_pixels : int = 512,
+        x_pixel_offset : int = 0,
+        y_pixel_offset : int = 0,
+        maximum_radius : float = 1.) -> Tensor:
+    """
+    The hexike polynomials up until `number_of_hexikes` on a square
+    array that `number_of_pixels` by `number_of_pixels`. The 
+    polynomials can be restricted to a smaller subset of the 
+    array by passing an explicit `maximum_radius`. The polynomial
+    will then be defined on the largest hexagon that fits with a 
+    circle of radius `maximum_radius`. 
+    
+    Parameters
+    ----------
+    number_of_hexikes : int = 15
+        The number of basis terms to generate. 
+    number_of_pixels : int = 512
+        The size of the array to compute the hexikes on.
+    x_pixel_offset : int
+        The offset of the aperture in the square output array in the 
+        x direction.
+    y_pixel_offset : int
+        The offset of the aperture in the square output array in the 
+        y direction. 
+    maximum_radius : float = 1.
+        The radius of the the smallest circle that can contain the 
+        hexagonal surface. 
+
+    Returns
+    -------
+    hexikes : Tensor
+        The hexike polynomials evaluated on the square arrays
+        containing the hexagonal apertures until `maximum_radius`.
+        The leading dimension is `number_of_hexikes` long and 
+        each stacked array is a basis term. The final shape is:
+        ```py
+        hexikes.shape == (number_of_hexikes, number_of_pixels, number_of_pixels)
+        ```
+    """
+    aperture = _hexagonal_aperture(number_of_pixels, x_pixel_offset,
+        y_pixel_offset, maximum_radius)
+
+    pixel_area = aperture.sum()
+    shape = (number_of_hexikes, number_of_pixels, number_of_pixels)
+    zernikes = dLux.utils.zernike_basis(number_of_hexikes, 
+        number_of_pixels)
+
+    # So the issue is currently that I cannot pass a negative offset
+    # The aperture is fine with these it is just that I cannot set 
+    # them properly in the following array.
+    if y_pixel_offset >= 0:
+        if x_pixel_offset >= 0:
+            offset_zernikes = np.zeros(shape)\
+                .at[:, : number_of_pixels - y_pixel_offset, 
+                    : number_of_pixels - x_pixel_offset]\
+                .set(zernikes[:, y_pixel_offset :, x_pixel_offset :])
+        else:
+            offset_zernikes = np.zeros(shape)\
+                .at[:, : number_of_pixels - y_pixel_offset, -x_pixel_offset :]\
+                .set(zernikes[:, y_pixel_offset :, : x_pixel_offset])
+    else:
+        if x_pixel_offset >= 0:
+            offset_zernikes = np.zeros(shape)\
+                .at[:, -y_pixel_offset :, : number_of_pixels - x_pixel_offset]\
+                .set(zernikes[:, : y_pixel_offset, x_pixel_offset :])
+        else:
+            offset_zernikes = np.zeros(shape)\
+                .at[:, -y_pixel_offset :, -x_pixel_offset :]\
+                .set(zernikes[:, : y_pixel_offset, : x_pixel_offset])
+
+
+
+    offset_hexikes = np.zeros(shape).at[0].set(aperture)
+    
+    for j in np.arange(1, number_of_hexikes): # Index of the zernike
+        intermediate = offset_zernikes[j + 1] * aperture
+
+        coefficients = -1 / pixel_area * \
+           ((offset_zernikes[j + 1] * offset_hexikes[1 : j + 1]) * aperture)\
+            .sum(axis = 0) 
+
+        intermediate += (coefficients * offset_hexikes[1 : j + 1])\
+            .sum(axis = 0)
+
+        offset_hexikes = offset_hexikes\
+            .at[j + 1]\
+            .set(intermediate / \
+                np.sqrt((intermediate ** 2).sum() / pixel_area))
+
+    return offset_hexikes
+
+
diff --git a/src/model_layers.py b/dLux/utils/models.py
similarity index 100%
rename from src/model_layers.py
rename to dLux/utils/models.py
diff --git a/src/plotting.py b/dLux/utils/plotting.py
similarity index 100%
rename from src/plotting.py
rename to dLux/utils/plotting.py
diff --git a/src/zernike.py b/dLux/utils/zernike.py
similarity index 96%
rename from src/zernike.py
rename to dLux/utils/zernike.py
index 08e50248..c5f69794 100644
--- a/src/zernike.py
+++ b/dLux/utils/zernike.py
@@ -121,7 +121,7 @@ def noll_indices(j):
         m = row_m[resid] * sign
     return n, m
 
-def zernike(n, m, npix=100, rho=None, theta=None, outside=np.nan,
+def zernike(n, m, npix=100, rho=None, theta=None, outside=0.,
             noll_normalize=True, **kwargs):
     """Return the Zernike polynomial Z[m,n] for a given pupil.
 
@@ -217,7 +217,7 @@ def R(n, m, rho):
     n, m : int
         Zernike function degree
     rho : array
-        Image plane radial coordinates. `rho` should be 1 at the desired pixel radius of the
+        Image plane radial coordnates. `rho` should be 1 at the desired pixel radius of the
         unit circle
     """
 
@@ -230,7 +230,8 @@ def R(n, m, rho):
         for k in range(int((n - m) / 2) + 1):
         
             coef = ((-1) ** k * factorial(n - k) /
-                    (factorial(k) * factorial( int((n + m) / 2) - k) * factorial(int((n - m) / 2) - k)))
+                    (factorial(k) * factorial(int((n + m) / 2) - k) *\
+                        factorial(int((n - m) / 2) - k)))
             output += coef * rho ** (n - 2 * k)
             
-        return output
\ No newline at end of file
+        return output
diff --git a/src/wavefronts.py b/dLux/wavefronts.py
similarity index 100%
rename from src/wavefronts.py
rename to dLux/wavefronts.py
diff --git a/docs/user_documentation/layers/apply_apperture.md b/docs/user_documentation/layers/apply_aperture.md
similarity index 100%
rename from docs/user_documentation/layers/apply_apperture.md
rename to docs/user_documentation/layers/apply_aperture.md
diff --git a/docs/wavefronts.html b/docs/wavefronts.html
new file mode 100644
index 00000000..a75c69e6
--- /dev/null
+++ b/docs/wavefronts.html
@@ -0,0 +1,12 @@
+<!--[if IE]><meta http-equiv="X-UA-Compatible" content="IE=5,IE=9" ><![endif]-->
+<!DOCTYPE html>
+<html>
+<head>
+<title>wavefronts.html</title>
+<meta charset="utf-8"/>
+</head>
+<body>
+<div class="mxgraph" style="max-width:100%;border:1px solid transparent;" data-mxgraph="{&quot;highlight&quot;:&quot;#0000ff&quot;,&quot;nav&quot;:true,&quot;resize&quot;:true,&quot;xml&quot;:&quot;&lt;mxfile host=\&quot;app.diagrams.net\&quot; modified=\&quot;2022-07-08T04:18:33.453Z\&quot; agent=\&quot;5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/103.0.0.0 Safari/537.36\&quot; etag=\&quot;M72DNKgpTv6vZjLLJZNZ\&quot; version=\&quot;20.0.4\&quot; type=\&quot;device\&quot;&gt;&lt;diagram id=\&quot;a0kQ9t1lwLzjGA3GEnq8\&quot; name=\&quot;Page-1\&quot;&gt;7Z1bj9o4FIB/DVL3oStyg/A40Nnuw1Sq2tVeniIPMeDWxJFjZqC/fu0k5hpTZobkIPVIowqMQ7DPF5+bfdoLJsv1R0nyxSeRUt7z++m6F3zo+X4Yxvpf07CpGqJBv2qYS5ZWTd6u4Sv7QetG223FUlocdFRCcMXyw8apyDI6VQdtRErxfNhtJvjhXXMypycNX6eEn7b+w1K1qFpjf7hr/5Oy+cLe2RuMqk+WxHauR1IsSCqe95qC+14wkUKo6tVyPaHczJ2dl+q6Pxyfbn+YpJm65IL7H9/S4f3Hf9f39CGb/zcV4vvf7+P6a54IX9Ujrn+t2tgpoKmekfqtkGoh5iIj/H7XOpZilaXU3Kev3+36PAiR60ZPN36jSm1q8ZKVErppoZa8/rS6p7mRc3B1UyFWckrPjCiqISFyTtWZfmGwlYFml4olVXKjL5SUE8WeDn8IqSmab/vtJlq/qOf6BfPuWSp+vYmPfdCJj0a/6sR71TdCzXx0Mu9jUtB/yBOdSaHHfCyE4pktOcnMbM/051YeZvIIZ/NMv57quaJSNzxRqZhetO/qD5SRwni6YDx9IBuxMtNSKDL9bt+NF0KyH/priRWI/liqWlxB/6DHV3NlLWpJC93nsxWTd9T0iawPOj6QQtUNU8E5yQv2uB3GUkuMZWOhlFjWneo50sOh6yP9Vl5N3au/Q7j11wRBWF2yOdSvzzul5lnVvNhTaEHfjUN9sy9a75Jsrof1krsF/Ya7jYbh4e0I1wLOiKJj88AVJxBuB/p6LgcN68HY/IRlzplapWZYveCusifI5oRTLSlV8iPFdzoRXGggP2SiApdxftRk2eV0ppzkFjmZsmz+UPb5EO5avtRzZZqEvnbGS5NiwdKUZuW6pIgiFWKGp1ywTJVzGY31n57ySf/3qBfpHz7R773de/1nuks1EZkeC2ElV1Tz+0wNw2409yF0P/mnZG4OCfsZG/6gLZUQOAjIF3qBQum3Kv1oAC193yH9Z62ZOM3m+mdZBPR8k1NFhQi8DYHYh0bg1DSpEBCzWaFNKiv+O1wBWhC/14+h5R+6FABbU54UZTQA14BWIbAGIxgEwwYGjoTMWSm8Ssg2GuS9SsJLLavSkaxF+lfpNr73TsQenIo9aBAxJ4+UfxYFU0yY75dV3yPRg0l3cKGNF7ck3NjxgCcJy5hKknf1Uz2xq/xv2+f9yEnV+AXezlPDJeCKkEDrgabIkKFkTlWypwrePQmW7gBBhdAKDSNoheA1RWgtDpVleEQCeohtkLBNoMCR0JQkMSQUOxLOaw1E4rpIhOBRg3O6Yhc5QFXRBQ1D8ACCy74sDmko5Y+LREdYjKDtSa/Jp7SLxDa9gEZEF0kGH9ycdOWZSu/CZBoQhC5AiKCtSd8VbSwOVgWCBmWHVAyhDUrflYQstssDEtEhEWEf2qi0AehGhSEp4agvuuDAh7YifVf0wXDAlmTOMiI3CEMXMETQVmTjtlEDQ7ZaPlKZiFkVqS6OeGCoIq5PQwxuSroczOXK7JDnm2N70vAwnNhIlMuWmPi9uz7icu2dTn1wG9O51TFN923MyyjBVGh7qATgxqdrS9QqT4miJS1CWlwm1uJAVrpnJYI2UAOXTfJs5Z8okeTFDG3ULniIwW1UV27MqBmRp6hkbgKUwaWHMtoDxZU2M57tVGjTla6TmZBLXDi64CGAtk8DVww00wxoeZ2kSDAE2ioQA2gr1JrBZ3bk6Y4yZeaEWPFOzwzqke4xicEN0HMR0goTU1wAtUgHNAw9aPMzcLmuOxryeqP8cZQUkWgFiRDa0AxcKXeWmeEm64RkaXKcQ0HrolUqBuDmpsv9sFQgD13yMAK3Nl3ZlJoHXB86PRXsgZuVrnQJM1VHcsG107EfAX8Ugu/oUKvcnBpFKq5LRQhtXoYuZ0OPjKwZ4ck+HnvnCssU/CknuIq0yssQ2vYMXdmR4uhUIR4O6JKLEbj16cqS5CRNlDgbzEIYrgrDyIc2PUNXoHMqRY40dEtDCG14hqd5ECyTd1mZvMsL4tmdlDaKfaGd4NlH9Y0F8RoONV9Uf6+DgnihK1qWa84StcmxJtpbVyBX8dlbKYkXXlAkFavhvFa8l9a8a6saTuRySbAczg0tAgPoSEf087Qq1sPpDIcRtE6IXD4K1sPplgTPg45dxKcofCSromAkQ0fl2o7KsetwqaeyPdP8Nk+l4Xa34qjErv2CWLr7WkuR6z9ruBVHJXb6qli7u33xgxfvjl17vbB4d1cMgFfvjl1peqze3YX84ct3x659O1i+uzsKwOt3x027+TBkeS35glfwdqXLMWZ5S8sAfA3vc1WbMWjZNQ/gVbxH54LYGLXsDgX4Mt52gzmW8b4ZJsDreI9cIUSs4w2AA3gh75ErnIiFvCG5AK/kPTp3phQreXebcwAv5T1yRRyxlHe3JIDX8vYairPeZfMVJxIT4ddOhEfWt3/pll3vVYnwC253K4lwr7EQKAY/L11nPBd5L0xztxX89CKXxsFDZK0qoJ+CAW2YepEz+YXHyDrHAXzLrhe5alrgQTIAHlrcs6vfSiHUvh0hSb74JFJqevwP&lt;/diagram&gt;&lt;/mxfile&gt;&quot;,&quot;toolbar&quot;:&quot;pages zoom layers lightbox&quot;,&quot;page&quot;:0}"></div>
+<script type="text/javascript" src="https://app.diagrams.net/js/viewer-static.min.js"></script>
+</body>
+</html>
diff --git a/setup.py b/setup.py
index 129b3270..9e2b0806 100644
--- a/setup.py
+++ b/setup.py
@@ -30,7 +30,7 @@ def find_version(*file_paths):
 
 setuptools.setup(
     name="dLux",
-    version=find_version("src", "__init__.py"),
+    # version=find_version("dLux", "__init__.py"),
     description="A fully differentiable optical simulator build in Jax",
     long_description=long_description,
     # long_description_content_type="text/markdown",
@@ -43,8 +43,8 @@ def find_version(*file_paths):
         "Bug Tracker": "https://github.com/LouisDesdoigts/dLux/issues",
     },
     
-    package_dir={"dLux": "src"},
-    packages=["dLux"],
+    # package_dir={"": "src"},
+    # packages=["dLux"],
     
     install_requires=install_requires,
     
@@ -52,4 +52,7 @@ def find_version(*file_paths):
         "Programming Language :: Python :: 3",
         "Operating System :: OS Independent",
     ],
-)
\ No newline at end of file
+
+    # packages = ["src"] + setuptools.find_namespace_packages(where = "src")
+    packages = ["dLux", "dLux/utils"]
+)
diff --git a/src/__init__.py b/src/__init__.py
deleted file mode 100644
index 537e1d81..00000000
--- a/src/__init__.py
+++ /dev/null
@@ -1,11 +0,0 @@
-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 *


Test patch:
diff --git a/dLux/dev/tests/layers.py b/dLux/dev/tests/layers.py
new file mode 100644
index 00000000..41b2ad8e
--- /dev/null
+++ b/dLux/dev/tests/layers.py
@@ -0,0 +1,15 @@
+from layers import *
+from matplotlib import pyplot
+from jax import numpy as np
+from jax.config import config
+from typing import TypeVar
+
+config.update("jax_debug_nans", True)
+config.update("jax_enable_x64", True)
+
+
+class TestHexagonalAperture(object):
+    """
+    Blank stub, since tests are best performed with visual comparison 
+    to poppy. Will likely be filled in future releases.
+    """