Apertures¤

class CircularAperture(DynamicAperture):
    """
    A dynamically generated circular aperture parameterised by its radius. Both jit and
    grad compatible.

    ??? abstract "UML"
        ![UML](../../assets/uml/CircularAperture.png)

    Attributes
    ----------
    radius: float, meters
        The radius of the aperture.
    transformation: CoordTransform
        The object that applies the coordinate transformations to the aperture.
    occulting: bool
        Is the aperture occulting or transmissive. False results in a
    softening: float, pixels
        The approximate pixel width of the soft boundary applied to the aperture.
    normalise : bool
        Whether to normalise the wavefront after passing through the aperture.
    extent : float, property
        Derived property from `radius`; maximum aperture extent.
    nsides : int, property
        Derived property from aperture type; `0` for circular symmetry.
    """

    radius: float

    def __init__(
        self: ApertureLayer,
        radius: float,
        transformation: CoordTransform = None,
        occulting: bool = False,
        softening: float = 1.0,
        normalise: bool = False,
    ):
        """
        Parameters
        ----------
        radius: float, meters
            The radius of the aperture.
        transformation: CoordTransform = None
            The object that applies the coordinate transformations to the aperture.
        occulting: bool = False
            Is the aperture occulting or transmissive. False results in a
            transmissive aperture, and True results in an occulting aperture.
        softening: float, pixels = 1.0
            The approximate pixel width of the soft boundary applied to the aperture.
        normalise : bool = False
            Whether to normalise the wavefront after passing through the aperture.
        """
        super().__init__(
            transformation=transformation,
            occulting=occulting,
            softening=softening,
            normalise=normalise,
        )
        self.radius = float(radius)

    def transmission(self: ApertureLayer, coords: Array, pixel_scale: float) -> Array:
        if self.transformation is not None:
            coords = self.transformation(coords)
        clip_val = pixel_scale * self.softness / 2
        return dlu.soft_circle(coords, self.radius, clip_val, self.occulting)

    @property
    def extent(self: ApertureLayer) -> float:
        return self.radius

    @property
    def nsides(self: ApertureLayer) -> int:
        return 0

class RectangularAperture(DynamicAperture):
    """
    A dynamically generated rectangular aperture parameterised by its width and height.
    Both jit and grad compatible.

    ??? abstract "UML"
        ![UML](../../assets/uml/RectangularAperture.png)

    Attributes
    ----------
    height: float, meters
        The length of the aperture in the y-direction.
    width: float, meters
        The length of the aperture in the x-direction.
    transformation: CoordTransform
        The object that applies the coordinate transformations to the aperture.
    occulting: bool
        Is the aperture occulting or transmissive. False results in a
        transmissive aperture, and True results in an occulting aperture.
    softening: float, pixels
        The approximate pixel width of the soft boundary applied to the aperture.
    normalise : bool
        Whether to normalise the wavefront after passing through the aperture.
    extent : float, property
        Derived property from `height` and `width`; half-diagonal extent.
    nsides : int, property
        Derived property from aperture type; `4` sides.
    """

    height: float
    width: float

    def __init__(
        self: ApertureLayer,
        height: float,
        width: float,
        transformation: CoordTransform = None,
        occulting: bool = False,
        softening: float = 1.0,
        normalise: bool = False,
    ):
        """
        Parameters
        ----------
        height: Array, meters
            The length of the aperture in the y-direction.
        width: Array, meters
            The length of the aperture in the x-direction.
        transformation: CoordTransform
            The object that applies the coordinate transformations to the aperture.
        occulting: bool = False
            Is the aperture occulting or transmissive. False results in a
            transmissive aperture, and True results in an occulting aperture.
        softening: float, pixels = 1.0
            The approximate pixel width of the soft boundary applied to the aperture.
        normalise : bool = False
            Whether to normalise the wavefront after passing through the aperture.
        """
        self.height = float(height)
        self.width = float(width)

        super().__init__(
            transformation=transformation,
            occulting=occulting,
            softening=softening,
            normalise=normalise,
        )

    def transmission(self: ApertureLayer, coords: Array, pixel_scale: float) -> Array:
        if self.transformation is not None:
            coords = self.transformation(coords)
        clip_val = pixel_scale * self.softness / 2
        return dlu.soft_rectangle(
            coords, self.width, self.height, clip_val, self.occulting
        )

    @property
    def extent(self: ApertureLayer) -> float:
        return np.hypot(self.height / 2.0, self.width / 2.0)

    @property
    def nsides(self: ApertureLayer) -> int:
        return 4

class RegPolyAperture(DynamicAperture):
    """
    Creates a dynamically generated regular polygon aperture parameterised by its
    number of sides and maximum radius. Both jit and grad compatible.

    ??? abstract "UML"
        ![UML](../../assets/uml/RegPolyAperture.png)

    Attributes
    ----------
    nsides: int
        The number of sides of the aperture.
    rmax: float, meters
        The maximum radius to the vertices from its center.
    transformation: CoordTransform
        The object that applies the coordinate transformations to the aperture.
    occulting: bool
        Is the aperture occulting or transmissive. False results in a
        transmissive aperture, and True results in an occulting aperture.
    softening: float, pixels
        The approximate pixel width of the soft boundary applied to the aperture.
    normalise : bool
        Whether to normalise the wavefront after passing through the aperture.
    extent : float, property
        Derived property from `rmax`; maximum aperture extent.
    """

    _nsides: int  # Enables nsides to be a class property to stick to design pattern
    rmax: float

    def __init__(
        self: ApertureLayer,
        nsides: int,
        rmax: float,
        transformation: CoordTransform = None,
        occulting: bool = False,
        softening: float = 1.0,
        normalise: bool = False,
    ):
        """
        Parameters
        ----------
        nsides: int
            The number of sides of the aperture.
        rmax: float, meters
            The maximum radius to the vertices from its center.
        transformation: CoordTransform
            The object that applies the coordinate transformations to the aperture.
        occulting: bool = False
            Is the aperture occulting or transmissive. False results in a
            transmissive aperture, and True results in an occulting aperture.
        softening: float, pixels = 1.0
            The approximate pixel width of the soft boundary applied to the aperture.
        normalise : bool = False
            Whether to normalise the wavefront after passing through the aperture.
        """
        self._nsides = int(nsides)
        self.rmax = float(rmax)

        super().__init__(
            transformation=transformation,
            occulting=occulting,
            softening=softening,
            normalise=normalise,
        )

    def transmission(self: ApertureLayer, coords: Array, pixel_scale: float) -> Array:
        if self.transformation is not None:
            coords = self.transformation(coords)
        clip_val = pixel_scale * self.softness / 2
        return dlu.soft_reg_polygon(
            coords, self.rmax, self.nsides, clip_val, self.occulting
        )

    @property
    def extent(self: ApertureLayer) -> float:
        return self.rmax

    @property
    def nsides(self: ApertureLayer) -> int:
        return self._nsides

class Spider(DynamicAperture):
    """
    Creates a dynamically generated spider aperture parameterised by its arm width and
    number of arms. Both jit and grad compatible.

    ??? abstract "UML"
        ![UML](../../assets/uml/Spider.png)

    Attributes
    ----------
    width: float, meters
        The width of the spider.
    angles: Array, degrees
        The angle of each arm of the spider.
    transformation: CoordTransform
        The object that applies the coordinate transformations to the aperture.
    occulting: bool
        Is the aperture occulting or transmissive. False results in a
        transmissive aperture, and True results in an occulting aperture.
    softening: float, pixels
        The approximate pixel width of the soft boundary applied to the aperture.
    normalise : bool
        Whether to normalise the wavefront after passing through the aperture.
    extent : float, property
        Derived property from aperture type; undefined for spiders and raises.
    nsides : int, property
        Derived property from aperture type; undefined for spiders and raises.
    """

    width: float
    angles: Array

    def __init__(
        self: ApertureLayer,
        width: float,
        angles: Array,
        transformation: CoordTransform = None,
        occulting: bool = True,
        softening: float = 1.0,
        normalise: bool = False,
    ):
        """
        Parameters
        ----------
        width: float, meters
            The width of the spider.
        angles: Array, degrees
            The angle of each arm of the spider.
        transformation: CoordTransform
            The object that applies the coordinate transformations to the aperture.
        occulting: bool = True
            Is the aperture occulting or transmissive. False results in a
            transmissive aperture, and True results in an occulting aperture.
        softening: float, pixels = 1.0
            The approximate pixel width of the soft boundary applied to the aperture.
        normalise : bool = False
            Whether to normalise the wavefront after passing through the aperture.
        """
        super().__init__(
            transformation=transformation,
            occulting=occulting,
            softening=softening,
            normalise=normalise,
        )

        self.width = float(width)
        self.angles = np.asarray(angles, dtype=float)

    def transmission(self: ApertureLayer, coords: Array, pixel_scale: float) -> Array:
        if self.transformation is not None:
            coords = self.transformation(coords)
        clip_val = pixel_scale * self.softness / 2
        return dlu.soft_spider(
            coords, self.width, self.angles, clip_val, self.occulting
        )

    @property
    def extent(self: ApertureLayer) -> float:
        raise TypeError("Spiders do not have an extent.")

    @property
    def nsides(self: ApertureLayer) -> int:
        raise TypeError("Spiders do not have a number of sides.")

class SquareAperture(DynamicAperture):
    """
    A dynamically generated square aperture parameterised by its side length.
    Both jit and grad compatible.

    ??? abstract "UML"
        ![UML](../../assets/uml/SquareAperture.png)

    Attributes
    ----------
    width: float, meters
        The side length of the aperture.
    transformation: CoordTransform
        The object that applies the coordinate transformations to the aperture.
    occulting: bool
        Is the aperture occulting or transmissive. False results in a
        transmissive aperture, and True results in an occulting aperture.
    softening: float, pixels
        The approximate pixel width of the soft boundary applied to the aperture.
    normalise : bool
        Whether to normalise the wavefront after passing through the aperture.
    extent : float, property
        Derived property from `width`; circumscribed half-diagonal scale.
    nsides : int, property
        Derived property from aperture type; `4` sides.
    """

    width: float

    def __init__(
        self: ApertureLayer,
        width: float,
        transformation: CoordTransform = None,
        occulting: bool = False,
        softening: float = 1.0,
        normalise: bool = False,
    ):
        """
        Parameters
        ----------
        width: Array, meters
            The side length of the aperture.
        transformation: CoordTransform
            The object that applies the coordinate transformations to the aperture.
        occulting: bool = False
            Is the aperture occulting or transmissive. False results in a
            transmissive aperture, and True results in an occulting aperture.
        softening: float, pixels = 1.0
            The approximate pixel width of the soft boundary applied to the aperture.
        normalise : bool = False
            Whether to normalise the wavefront after passing through the aperture.
        """
        super().__init__(
            transformation=transformation,
            occulting=occulting,
            softening=softening,
            normalise=normalise,
        )

        self.width = float(width)

    def transmission(self: ApertureLayer, coords: Array, pixel_scale: float) -> Array:
        if self.transformation is not None:
            coords = self.transformation(coords)
        clip_val = pixel_scale * self.softness / 2
        return dlu.soft_square(coords, self.width, clip_val, self.occulting)

    @property
    def extent(self: ApertureLayer) -> float:
        return np.sqrt(2) * self.width

    @property
    def nsides(self: ApertureLayer) -> int:
        return 4

class AberratedAperture(BasisLayer, ApertureLayer):
    """
    Creates a dynamically generated Aperture with aberrations. Both jit and grad
    compatible.

    ??? abstract "UML"
        ![UML](../../assets/uml/AberratedAperture.png)

    Attributes
    ----------
    aperture: ApertureLayer
        The aperture on which the aberration basis is defined.
    basis: list[Zernike]
        A list of basis functions that generate the basis vectors.
    coefficients: Array
        The amplitude of each basis vector of the aberrations.
    as_phase : bool
        Whether to apply the basis as a phase or OPD. If True the basis is
        applied as a phase, else it is applied as an OPD.
    normalise : bool
        Whether to normalise the wavefront after passing through the aperture.
    """

    aperture: ApertureLayer

    def __init__(
        self: ApertureLayer,
        aperture: ApertureLayer,
        noll_inds: Array,
        coefficients: Array = None,
        as_phase: bool = False,
    ):
        """
        Parameters
        ----------
        aperture: ApertureLayer
            The aperture on which the aberration basis is defined.
        noll_inds: Array
            The noll indices of the basis functions to use.
        coefficients: Array = None
            The amplitude of each basis vector of the aberrations.
        as_phase : bool = False
            Whether to apply the basis as a phase or OPD. If True the basis is
            applied as a phase, else it is applied as an OPD.
        """
        # Ensure aperture is dynamic
        if not isinstance(aperture, DynamicAperture):
            raise TypeError(
                "AberratedApertures cannot contain Static, "
                "Compound or Multi Apertures. AberratedApertures can be "
                "placed in Compound or Multi Apertures, which can then be "
                "promoted to Static."
            )

        super().__init__(normalise=aperture.normalise, as_phase=as_phase)

        # Ensure transmissive
        if aperture.occulting:
            raise TypeError("AberratedApertures cannot be occulting.")

        if isinstance(aperture, Spider):
            raise TypeError("AberratedApertures cannot be spiders.")

        # Set Aperture
        self.aperture = aperture
        self.basis = ZernikeBasis(noll_inds)

        if coefficients is None:
            coefficients = np.zeros(len(noll_inds))
        self.coefficients = np.asarray(coefficients, dtype=float)

    def calculate(self: ApertureLayer):
        """
        Required abstract method. Raises NotImplementedError as it is invalid here.
        """
        raise NotImplementedError(
            "Aberrated Apertures cannot use the .calculate() method because "
            "they need coords to be generated on. please use "
            ".eval_basis(coords) method instead."
        )

    def transmission(self: ApertureLayer, coords: Array, pixel_scale: float) -> Array:
        return self.aperture.transmission(coords, pixel_scale)

    def calc_basis(self: ApertureLayer, coords: Array) -> Array:
        """
        Calculates the basis vectors at the given coordinates.

        Parameters
        ----------
        coords : Array
            The coordinates to calculate the basis vectors on.

        Returns
        -------
        basis : Array
            The basis vectors at the given coordinates.
        """
        if self.aperture.transformation is not None:
            coords = self.aperture.transformation(coords)
        coords /= self.aperture.extent
        return self.basis.calculate_basis(coords, self.aperture.nsides)

    def eval_basis(self: ApertureLayer, coords: Array) -> Array:
        """
        Evaluates the basis vectors at the given coordinates.

        Parameters
        ----------
        coords : Array
            The coordinates to evaluate the basis vectors on.

        Returns
        -------
        aberrations : Array
            The aberrations at the given coordinates.
        """
        basis = self.calc_basis(coords)
        return dlu.eval_basis(basis, self.coefficients)

    def __call__(self: AberratedAperture, wavefront: Wavefront) -> Wavefront:
        """
        Applies the aperture transmission and aberrations to the wavefront.

        Parameters
        ----------
        wavefront : Wavefront
            The wavefront to operate on.

        Returns
        -------
        wavefront : Wavefront
            The transformed wavefront.
        """
        # Transmission
        wavefront *= self.transmission(wavefront.coordinates(), wavefront.pixel_scale)
        if self.normalise:
            wavefront = wavefront.normalise()

        # Transform coordinate
        if self.aperture.transformation is not None:
            coords = self.aperture.transformation(wavefront.coordinates())
        else:
            coords = wavefront.coordinates()

        # Aberrations
        aberrations = self.eval_basis(coords)
        if self.as_phase:
            wavefront = wavefront.add_phase(aberrations)
        else:
            wavefront = wavefront.add_opd(aberrations)
        return wavefront

class CompoundAperture(CompositeAperture):
    """
    Dynamically generates an aperture from a series of overlapping sub-apertures. Both
    jit and grad compatible.

    This class combines the aperture via a _multiplication_ of the sub-apertures. An
    example would be a HST-like aperture with an obscuring secondary mirror and spiders.

    This class can only contain a _single_ AberratedAperture.

    If you want to combine apertures via an _addition_ of the sub-apertures such as an
    aperture mask, use the MultiAperture class.

    Note that this class cannot contain a MultiAperture, but MultiApertures can
    contain CompoundApertures.

    ??? abstract "UML"
        ![UML](../../assets/uml/CompoundAperture.png)

    Attributes
    ----------
    apertures: dict
        The sub-apertures that make up the full aperture.
    transformation: CoordTransform
        The object that applies the coordinate transformations to the aperture.
    normalise : bool
        Whether to normalise the wavefront after passing through the aperture.
    """

    def __init__(
        self: ApertureLayer,
        apertures: list,
        transformation: CoordTransform = None,
        normalise: bool = False,
    ):
        """
        Parameters
        ----------
        apertures: list[ApertureLayer]
            The sub-apertures that make up the full aperture.
        transformation: CoordTransform
            The object that applies the coordinate transformations to the aperture.
        normalise : bool = False
            Whether to normalise the wavefront after passing through the aperture.
        """
        super().__init__(
            apertures=apertures,
            transformation=transformation,
            normalise=normalise,
        )

        # Check for more than one aberrated aperture
        naberrated = 0
        for aperture in self.apertures.values():
            if isinstance(aperture, AberratedAperture):
                naberrated += 1
            if naberrated > 1:
                raise TypeError(
                    "CompoundAperture can only have a single " "AberratedAperture."
                )

    def transmission(self: ApertureLayer, coords: Array, pixel_scale: float) -> Array:
        if self.transformation is not None:
            coords = self.transformation(coords)
        return self.transmissions(coords, pixel_scale).prod(0)

class MultiAperture(CompositeAperture):
    """
    Dynamically generates an aperture from a series of separated sub-apertures. Both
    jit and grad compatible.

    This class combines the aperture via an _addition_ of the sub-apertures. An example
    would be an aperture mask with multiple holes.

    This class can only contain a _multiple_ AberratedAperture, or CompoundApertures.

    If you want to combine apertures via a _multiplication_ of the sub-apertures such
    as HST-like aperture, use the CompoundAperture class.

    ??? abstract "UML"
        ![UML](../../assets/uml/MultiAperture.png)

    Attributes
    ----------
    apertures: dict
        The sub-apertures that make up the full aperture.
    transformation: CoordTransform
        The object that applies the coordinate transformations to the aperture.
    normalise : bool
        Whether to normalise the wavefront after passing through the aperture.
    """

    def eval_basis(self: ApertureLayer, coords: Array) -> Array:
        """
        Evaluates the basis vectors at the given coordinates.

        Parameters
        ----------
        coords : Array
            The coordinates to evaluate the basis vectors on.

        Returns
        -------
        aberrations : Array
            The aberrations at the given coordinates.
        """
        return super().eval_basis(coords).sum(0)

    def transmission(self: ApertureLayer, coords: Array, pixel_scale: float) -> Array:
        if self.transformation is not None:
            coords = self.transformation(coords)
        return self.transmissions(coords, pixel_scale).sum(0)