grater_jax.disk_model.winnie_jwst_fm

winnie_jwst_fm.py

JWST PSF grid and convolution utilities.

This module provides functions adapted from Kellen Lawson’s Winnie package for generating and applying JWST NIRCam coronagraphic PSFs. It includes routines for creating spatially varying PSF grids, convolving model images with PSFs, handling coronagraph masks, and performing image rotations/padding on CPU or GPU. GPU acceleration via CuPy is supported where available. This file is only used for its generate_nircam_grid function. This could be deprecated as more integrated support for Winnie could be developed.

Main functions

generate_nircam_psf_grid : Build a grid of NIRCam PSFs and index maps.
convolve_with_spatial_psfs : Apply spatially varying PSF convolution.
rotate_hypercube, pad_and_rotate_hypercube : Rotate image cubes.
generate_jwst_mask_image : Generate coronagraph mask transmission maps.
GPU-enabled helpers (rotate_image_gpu, psf_convolve_gpu, etc.).

Functions

`c_to_c_osamp`(center, osamp)	Convert a Cartesian coordinate to the corresponding oversampled value.
`convolve_with_spatial_psfs`(im0, psfs, psf_inds)	Convolve an image with a spatially varying PSF map.
`dist_to_pt`(pt[, nx, ny, dtype])	Compute a distance map to a point.
`free_gpu`(*args)	Free intermediate GPU buffers.
`generate_jwst_mask_image`(inst, cent[, ...])	Generate a JWST coronagraph mask transmission image.
`generate_nircam_psf_grid`(inst, coron_cents)	Generate a grid of synthetic NIRCam PSFs and associated index/throughput maps.
`isNone`(arg)	Convenience: robust None check that works with numpy arrays.
`mpl_stcen_extent`(im[, cent, pixelscale])	Compute a Matplotlib `extent` centered on a given star position.
`pad_and_rotate_hypercube`(hcube, angles[, ...])	Pad images to avoid pixel loss, then rotate a hypercube.
`pad_or_crop_about_pos`(im, pos, new_size[, ...])	Pad or crop an image about a given position.
`pad_or_crop_image`(im, new_size[, cent, ...])	Pad or crop an image about a specified center.
`propagate_nans_in_spatial_operation`(a, fn[, ...])	Apply a spatial operation while propagating NaNs (and optionally zeros).
`psf_convolve_cpu`(im, psf_im)	Convolve on CPU using SciPy's FFT-based convolution.
`psf_convolve_gpu`(im, psf_im)	Convolve on GPU using CuPy's FFT-based convolution.
`raw_model_to_convolved_model_cube`(input_im, ...)	Create a PSF-convolved sequence of model images at given parallactic angles.
`rotate_about_pos`(im, pos, angle[, new_cent, ...])	Rotate about an arbitrary position on CPU.
`rotate_about_pos_gpu`(im, pos, angle[, ...])	Rotate about an arbitrary position on GPU.
`rotate_hypercube`(hcube, angles[, cent, ...])	Rotate an array of images where the last two axes are `(y, x)`.
`rotate_image_cpu`(im, angle[, cent, ...])	Rotate an image by `angle` (degrees) using CPU operations.
`rotate_image_gpu`(im0, angle[, cent, ...])	Rotate an image by `angle` (degrees) using GPU operations (CuPy).
`xy_polar_ang_displacement`(x, y, dtheta)	Rotate Cartesian coordinates by an angle about the origin.
`xy_polar_ang_displacement_gpu`(x, y, dtheta)	GPU version of `xy_polar_ang_displacement()`.

grater_jax.disk_model.winnie_jwst_fm.raw_model_to_convolved_model_cube(input_im, parangs, psfs, psf_inds, im_mask=None, cent=None, use_gpu=False, ncores=-2)

Create a PSF-convolved sequence of model images at given parallactic angles.

Parameters:

input_im (ndarray, shape (ny, nx)) – Raw model image (oversampled if the PSFs are oversampled).
parangs (ndarray, shape (n_T,)) – Parallactic angles in degrees. Each output slice is rotated by the corresponding angle.
psfs (ndarray, shape (n_psf, ny, nx)) – Grid of spatially sampled PSFs used after rotation. If these PSFs are oversampled, input_im should be oversampled by the same factor.
psf_inds (ndarray, shape (ny, nx) or (n_T, ny, nx)) – If 2-D: index map selecting the PSF slice in psfs for each pixel. If 3-D: a different index map for each roll (matching parangs).
im_mask (ndarray, optional) – Coronagraph transmission map(s) multiplied into input_im before convolution. Shape as input_im (or stacked per roll if 3-D).
cent (array_like, optional) – Cartesian pixel coordinate [x, y] of the star in input_im.
use_gpu (bool, default False) – Use GPU (CuPy) for rotation/convolution where applicable.
ncores (int, default -2) – CPU threads for rotation when use_gpu=False; -2 uses all but one available core.

Returns:

imcube – For each angle in parangs, the rotated model convolved with the appropriate PSF.

Return type:

ndarray, shape (n_T, ny, nx)

grater_jax.disk_model.winnie_jwst_fm.convolve_with_spatial_psfs(im0, psfs, psf_inds, im_mask=None, use_gpu=False)

Convolve an image with a spatially varying PSF map.

Each pixel of im0 is convolved with the nearest PSF slice from psfs, as indicated by psf_inds.

Parameters:

im0 (ndarray, shape (ny, nx)) – Input image.
psfs (ndarray, shape (n_psf, ny, nx)) – Stack of PSF images.
psf_inds (ndarray, shape (ny, nx)) – Map of indices into psfs selecting the PSF per pixel.
im_mask (ndarray, optional) – Coronagraph throughput map multiplied into im0 before convolution.
use_gpu (bool, default False) – Use CuPy-based convolution.

Returns:

imcon – Convolved image.

Return type:

ndarray, shape (ny, nx)

grater_jax.disk_model.winnie_jwst_fm.psf_convolve_gpu(im, psf_im)

Convolve on GPU using CuPy’s FFT-based convolution.

Parameters:

im (ndarray) – Image to convolve.
psf_im (ndarray) – PSF kernel.

Returns:

Convolved image, same shape as im.

Return type:

ndarray

grater_jax.disk_model.winnie_jwst_fm.psf_convolve_cpu(im, psf_im)

Convolve on CPU using SciPy’s FFT-based convolution.

Parameters:

im (ndarray) – Image to convolve.
psf_im (ndarray) – PSF kernel.

Returns:

Convolved image, same shape as im.

Return type:

ndarray

grater_jax.disk_model.winnie_jwst_fm.c_to_c_osamp(center, osamp)

Convert a Cartesian coordinate to the corresponding oversampled value.

This uses typical zero-based indexing conventions.

Parameters:

center (array_like) – Coordinate(s) to convert.
osamp (float) – Oversampling factor.

Returns:

Converted coordinate(s).

Return type:

ndarray

grater_jax.disk_model.winnie_jwst_fm.generate_nircam_psf_grid(inst, coron_cents, source_spectrum=None, normalize=True, nr=6, ntheta=4, log_rscale=True, rmax=4, use_coeff=True)

Generate a grid of synthetic NIRCam PSFs and associated index/throughput maps.

Parameters:

inst (webbpsf_ext.webbpsf_ext_core.NIRCam_ext) – Configured NIRCam instrument object.
coron_cents (ndarray, shape (n_rolls, 2)) – Detector coordinates [x, y] of the coronagraph center in each roll after registration.
source_spectrum (pysynphot.Spectrum, optional) – Source spectrum used for PSF generation.
normalize (bool, default True) – Normalize each occulted PSF to unit sum. If model images are multiplied by the coronagraph throughput map before convolution, the throughput is recovered at higher resolution.
nr (int, default 6) – Number of radial PSF samples (excluding the origin).
ntheta (int, default 4) – Number of azimuthal PSF samples.
log_rscale (bool, default True) – If True, sample radii logarithmically between 0.01 and rmax arcsec.
rmax (float, default 4) – Maximum separation (arcsec) from the coronagraph center.
use_coeff (bool, default True) – Use coefficient-based PSF generation when available (inst.calc_psf_from_coeff).

Returns:

psfs (ndarray, shape (n_samples, ny, nx)) – Stack of sampled PSFs.
psf_inds_rolls (ndarray, shape (n_rolls, ny_os, nx_os)) – Map of nearest-PSF indices per (oversampled) detector pixel for each roll.
im_mask_rolls (ndarray, shape (n_rolls, ny_os, nx_os)) – Coronagraph transmission maps per roll (oversampled).
psf_offsets (ndarray, shape (2, n_samples)) – Offsets (x, y in arcsec) of each PSF sample.

grater_jax.disk_model.winnie_jwst_fm.generate_jwst_mask_image(inst, cent, return_oversample=True)

Generate a JWST coronagraph mask transmission image.

Parameters:

inst (webbpsf_ext.webbpsf_ext_core.NIRCam_ext) – Instrument object.
cent (array_like) – Center position [x, y] in oversampled pixels.
return_oversample (bool, default True) – If True, return the oversampled mask image. If False, rebin to the detector sampling.

Returns:

Transmission map (oversampled if requested).

Return type:

ndarray

grater_jax.disk_model.winnie_jwst_fm.propagate_nans_in_spatial_operation(a, fn, fn_args=None, fn_kwargs=None, fn_nan_kwargs=None, fn_zero_kwargs=None, prop_threshold=0, prop_zeros=True)

Apply a spatial operation while propagating NaNs (and optionally zeros).

The operation is intentionally liberal with propagation: after rotation or resampling, more pixels may be marked as NaN/zero depending on prop_threshold.

Parameters:

a (ndarray) – Input array.
fn (callable) – Spatial operation (e.g., scipy.ndimage.rotate).
fn_args (list, optional) – Positional arguments for fn.
fn_kwargs (dict, optional) – Keyword arguments for fn when applied to data.
fn_nan_kwargs (dict, optional) – Keyword arguments for fn when applied to the NaN mask.
fn_zero_kwargs (dict, optional) – Keyword arguments for fn when applied to the zero mask.
prop_threshold (float, default 0) – Mask values > prop_threshold are treated as True after transform.
prop_zeros (bool, default True) – Also propagate zeros as a mask.

Returns:

a_out – Transformed array with NaNs/zeros propagated through the operation.

Return type:

ndarray

grater_jax.disk_model.winnie_jwst_fm.pad_or_crop_image(im, new_size, cent=None, new_cent=None, cval=nan, prop_threshold=1e-06, order=3, mode='constant', prefilter=True)

Pad or crop an image about a specified center.

Parameters:

im (ndarray) – Image to transform.
new_size (tuple(int, int)) – New (ny, nx) shape.
cent (array_like, optional) – Current center [x, y]. If None, use geometric center.
new_cent (array_like, optional) – Desired center [x, y] in the output.
cval (float, default numpy.nan) – Fill value.
prop_threshold (float, default 1e-6) – Threshold used when propagating masks.
order (int, default 3) – Spline interpolation order for map_coordinates.
mode (str, default "constant") – Boundary mode for map_coordinates.
prefilter (bool, default True) – Prefilter for spline interpolation.

Returns:

Padded/cropped image.

Return type:

ndarray

grater_jax.disk_model.winnie_jwst_fm.pad_or_crop_about_pos(im, pos, new_size, new_cent=None, cval=nan, order=3, mode='constant', prefilter=True)

Pad or crop an image about a given position.

Parameters:

im (ndarray) – Input image or stack (…, ny, nx).
pos (array_like) – Original center [x, y].
new_size (tuple(int, int)) – Output size (ny, nx).
new_cent (array_like, optional) – New center [x, y] in the output.
cval (float, default numpy.nan) – Fill value.
order (int, default 3) – Spline interpolation order.
mode (str, default "constant") – Boundary mode for map_coordinates.
prefilter (bool, default True) – Prefilter for spline interpolation.

Returns:

Padded/cropped image (same ndim as input).

Return type:

ndarray

grater_jax.disk_model.winnie_jwst_fm.dist_to_pt(pt, nx=201, ny=201, dtype=<class 'float'>)

Compute a distance map to a point.

Parameters:

pt (array_like) – Point [x, y].
nx (int, optional) – Output size. Defaults to 201.
ny (int, optional) – Output size. Defaults to 201.
dtype (data-type, optional) – Output dtype. Default is float.

Returns:

Euclidean distance map.

Return type:

ndarray, shape (ny, nx)

grater_jax.disk_model.winnie_jwst_fm.rotate_image_cpu(im, angle, cent=None, new_cent=None, cval0=nan, prop_threshold=1e-06)

Rotate an image by angle (degrees) using CPU operations.

Exact zeros are treated similarly to NaNs and can be propagated.

Parameters:

im (ndarray) – Input image or stack (…, ny, nx).
angle (float) – Rotation angle in degrees.
cent (array_like, optional) – Center of rotation [x, y]. If None, use the geometric center.
new_cent (array_like, optional) – New desired center in the output.
cval0 (float, default numpy.nan) – Fill value.
prop_threshold (float, default 1e-6) – Threshold for mask propagation.

Returns:

Rotated image/stack.

Return type:

ndarray

grater_jax.disk_model.winnie_jwst_fm.rotate_image_gpu(im0, angle, cent=None, new_cent=None, cval0=nan, prop_threshold=1e-06)

Rotate an image by angle (degrees) using GPU operations (CuPy).

Parameters:

im0 (ndarray) – Input image.
angle (float) – Rotation angle in degrees.
cent (array_like, optional) – Center of rotation [x, y]. If None, use the geometric center.
new_cent (array_like, optional) – New desired center in the output.
cval0 (float, default numpy.nan) – Fill value.
prop_threshold (float, default 1e-6) – Threshold for mask propagation.

Returns:

Rotated image.

Return type:

ndarray

grater_jax.disk_model.winnie_jwst_fm.rotate_about_pos_gpu(im, pos, angle, new_cent=None, cval=nan, order=3, mode='constant', prefilter=True)

Rotate about an arbitrary position on GPU.

Parameters:

im (cupy.ndarray) – Input image or stack.
pos (array_like) – Center of rotation [x, y].
angle (float) – Rotation angle in degrees.
new_cent (array_like, optional) – New desired center in the output.
cval (float, default numpy.nan) – Fill value.
order (int, default 3) – Spline interpolation order.
mode (str, default "constant") – Boundary mode for map_coordinates.
prefilter (bool, default True) – Prefilter for spline interpolation.

Returns:

Rotated image/stack.

Return type:

cupy.ndarray

grater_jax.disk_model.winnie_jwst_fm.rotate_about_pos(im, pos, angle, new_cent=None, cval=nan, order=3, mode='constant', prefilter=True)

Rotate about an arbitrary position on CPU.

Parameters:

im (ndarray) – Input image or stack (…, ny, nx).
pos (array_like) – Center of rotation [x, y].
angle (float) – Rotation angle in degrees.
new_cent (array_like, optional) – New desired center in the output.
cval (float, default numpy.nan) – Fill value.
order (int, default 3) – Spline interpolation order.
mode (str, default "constant") – Boundary mode for map_coordinates.
prefilter (bool, default True) – Prefilter for spline interpolation.

Returns:

Rotated image/stack.

Return type:

ndarray

grater_jax.disk_model.winnie_jwst_fm.rotate_hypercube(hcube, angles, cent=None, new_cent=None, ncores=-2, use_gpu=False, cval0=0.0)

Rotate an array of images where the last two axes are (y, x).

Parameters:

hcube (ndarray) – Hypercube of images. For a time series of images with shape (ny, nx), this should have shape (n_T, ny, nx). For a time series of IFS cubes with n_L wavelengths, shape (n_T, n_L, ny, nx).
angles (ndarray, shape (n_T,)) – Rotation angles in degrees for the first dimension of hcube.
cent (array_like, optional) – Center of rotation [x, y].
new_cent (array_like, optional) – Desired center in the output.
ncores (int, default -2) – CPU threads for rotation when use_gpu=False; -2 uses all but one.
use_gpu (bool, default False) – Use CuPy-based rotation.
cval0 (float, default 0.0) – Fill value.

Returns:

Rotated hypercube with the same shape as hcube.

Return type:

ndarray

grater_jax.disk_model.winnie_jwst_fm.pad_and_rotate_hypercube(hcube, angles, cent=None, ncores=-2, use_gpu=False, cval0=nan)

Pad images to avoid pixel loss, then rotate a hypercube.

Parameters:

hcube (ndarray) – Hypercube to rotate.
angles (ndarray, shape (n_T,)) – Rotation angles in degrees.
cent (array_like, optional) – Center of rotation [x, y]. If None, uses geometric center.
ncores (int, default -2) – CPU threads for rotation when use_gpu=False.
use_gpu (bool, default False) – Use CuPy-based rotation.
cval0 (float, default numpy.nan) – Fill value.

Returns:

hcube_pad_rot (ndarray) – Padded and rotated hypercube.
cent_pad (ndarray) – New center after padding.

grater_jax.disk_model.winnie_jwst_fm.xy_polar_ang_displacement(x, y, dtheta)

Rotate Cartesian coordinates by an angle about the origin.

Parameters:

x (array_like) – Coordinates to rotate.
y (array_like) – Coordinates to rotate.
dtheta (float) – Rotation angle in degrees.

Returns:

newx, newy – Rotated coordinates.

Return type:

ndarray

grater_jax.disk_model.winnie_jwst_fm.xy_polar_ang_displacement_gpu(x, y, dtheta)

GPU version of xy_polar_ang_displacement().

Parameters:

x (cupy.ndarray) – Coordinates to rotate.
y (cupy.ndarray) – Coordinates to rotate.
dtheta (float) – Rotation angle in degrees.

Returns:

Rotated coordinates (newx, newy).

Return type:

cupy.ndarray

grater_jax.disk_model.winnie_jwst_fm.free_gpu(*args)

Free intermediate GPU buffers.

Parameters:: *args – Arrays to dereference before clearing the CuPy memory pools.
Returns:: Returns the list of inputs replaced by None (or None if only one).
Return type:: list or None

grater_jax.disk_model.winnie_jwst_fm.isNone(arg)

Convenience: robust None check that works with numpy arrays.

Parameters:: arg (Any) – Object to test.
Returns:: True if arg is None.
Return type:: bool

grater_jax.disk_model.winnie_jwst_fm.mpl_stcen_extent(im, cent=None, pixelscale=None)

Compute a Matplotlib extent centered on a given star position.

Parameters:

im (ndarray) – Image.
cent (array_like, optional) – Center coordinate [x, y]. If None, use geometric center.
pixelscale (float, optional) – Pixel scale to convert to physical units.

Returns:

Extent array: [xmin, xmax, ymin, ymax].

Return type:

ndarray, shape (4,)