yippy.offjax

Module for handling off-axis PSFs using JAX.

Classes

OffJAX

Class for handling off-axis PSFs using JAX.

Module Contents

class yippy.offjax.OffJAX(yip_dir, offax_data_file, offax_offsets_file, pixel_scale_arcsec, x_symmetric, y_symmetric, downsample_shape=None)

Bases: yippy.offax.OffAx

Class for handling off-axis PSFs using JAX.

This class inherits from OffAx and uses JAX for optimized computation. Memory-efficient: stores PSFs in flat array with index mapping.

Attributes:

pixel_scale_arcsec (Quantity):

Pixel scale of the PSF data in lambda/D.

center_x (Quantity):

Central x position in the PSF data.

center_y (Quantity):

Central y position in the PSF data.

flat_psfs (jnp.ndarray):

Flat array of PSF data with shape (N_psfs, H, W), cast as JAX array.

offset_to_flat_idx (jnp.ndarray):

2D index mapping from (x_idx, y_idx) -> flat_psfs index.

Parameters:

• yip_dir (pathlib.Path)

• offax_data_file (str)

• offax_offsets_file (str)

• pixel_scale_arcsec (astropy.units.Quantity)

• x_symmetric (bool)

• y_symmetric (bool)

• downsample_shape (tuple[int, int] | None)

flat_psfs

flat_x_offsets

flat_y_offsets

x_offsets

y_offsets

offset_to_flat_idx

kx

ky

create_psfs_kernel

create_psfs_j = None

create_psf_kernel_single = None

create_psfs

Create and return the PSFs at the specified off-axis positions.

Deprecated since version The: pure-Python implementation is deprecated. Use Coronagraph (which uses OffJAX) instead.

create_psf

Create and return the PSF at the specified off-axis position.

Deprecated since version The: pure-Python implementation is deprecated. Use Coronagraph (which uses OffJAX) instead.

Interpolates and returns the Point Spread Function (PSF) at the specified off-axis position (x, y). If the exact (x, y) position matches one of the PSFs in the YIP, that PSF is returned directly. Otherwise, the PSFs from the surrounding positions are combined using Gaussian weighting and Fourier interpolation to produce an interpolated PSF.

Args:

x (float):

The x-coordinate of the off-axis position.

y (float):

The y-coordinate of the off-axis position.

Returns:

np.ndarray:

The interpolated PSF corresponding to the input (x, y) position.

Notes:

• If self.type is “1d”, the (x, y) position is converted to a

radial separation and angle for interpolation. - If self.type is “2dq”, the (x, y) position is mirrored to the first quadrant, and the PSF is flipped accordingly after interpolation. - Gaussian weighting is used to combine the nearest PSFs when the exact (x, y) position does not match any precomputed PSF. The weighting is based on the distance from the input position. - The PSFs are shifted to align with the input position before combining, and the final PSF is normalized by the cumulative weight for each pixel.

create_psfs_parallel(x_vals, y_vals)

Create off-axis PSFs at multiple positions in parallel using shard_map.

Requires that hwoutils.set_host_device_count(N) was called at program startup to expose N CPU devices to JAX.