LouisDesdoigts / dLux

Differentiable optical models as parameterised neural networks in Jax using Zodiax
https://louisdesdoigts.github.io/dLux/
BSD 3-Clause "New" or "Revised" License
49 stars 6 forks source link

Update `ApplyJitter` class #249

Open maxecharles opened 1 year ago

maxecharles commented 1 year ago

Update the generate_kernel method of the existing class. Perhaps add in the ApplyLinearJitter class I wrote for Toliman work.

e.g.

class ApplyJitter(...):

...

  def generate_kernel(self: DetectorLayer, pixel_scale: Array) -> Array:
          """
          Generates the normalised Gaussian kernel.

          Returns
          -------
          kernel : Array
              The Gaussian kernel.
          """

          extent = self.kernel_size * pixel_scale
          x = np.linspace(0, extent, self.kernel_size) - 0.5*extent
          kernel = norm.pdf(x, scale=self.sigma) * norm.pdf(x[:, None], scale=self.sigma)
          return kernel / np.sum(kernel)

and turn this func into a DetectorLayer

def jitter_psf(rad: float, angle: float = 0, centre: tuple = (0,0), npsf: int = 10):
    """
    Returns a jittered PSF by summing a number of shifted PSFs.

    Parameters
    ----------
    rad : float
        The radius of the jitter in pixels.
    angle : float, optional
        The angle of the jitter in degrees, by default 0
    centre : tuple, optional
        The centre of the jitter in pixels, by default (0,0)
    npsf : int, optional
        The number of PSFs to sum, by default 10

    Returns
    -------
    np.ndarray
        The jittered PSF.
    """

    angle = np.deg2rad(angle) # converting to radius

    # converting to cartesian coordinates
    x = rad/2 * np.cos(angle)
    y = rad/2 * np.sin(angle)
    xs = np.linspace(-x, x, npsf)  # pixels
    ys = np.linspace( -y, y, npsf)  # pixels
    positions = pscale * (np.stack([xs, ys], axis=1) + np.array(centre))  # arcseconds

    psfs = vmap_prop(optics, source, positions)
    jit_psf = psfs.sum(0) / npsf  # adding and renormalising

    return jit_psf
maxecharles commented 1 year ago

Have created branch jitter (#253) and added the ApplyAsymmetricJitter class I wrote for TOLIMAN jitter fisher analysis.