Spatial method: grid, blending, convolution#

Spatially-varying PSF convolution: sample grid, blend weights and convolution engines.

The spatial method computes the PSF at a set of sample positions on the image (psf.spatial_psf_calculation builds this “model”) and gives every image pixel a blend of up to four of those samples. This module holds the parts that do not need the optics:

* bandlimited_weights: interpolation weights for band-limited functions of
  the source position (used by psf for mode 'rotate');
* positions_to_polar, equal_error_grid, make_blend_weights, blend_map: the
  sample grid (neighbouring samples differ in PSF shape by about tol) and
  the per-pixel blend weights;
* convolve_with_blended_psfs: each pixel convolved with its own blended PSF
  (the engine behind spatial_psf);
* make_psf_inds, convolve_with_spatial_psfs: the nearest-sample engine (one
  PSF per pixel, no blending), kept for direct use (tests/stage1_engine.py);
* spatial_iwa, spatial_psf, image_error_estimate: the final image of a
  target and the estimate of its error.

Conventions: images are (img_pixel, img_pixel) arrays indexed (row, col); the star sits at pixel coordinate c = (img_pixel-1)/2 on both axes; sample positions are (row, col) in pixels (not necessarily integers); PSF kernels have odd sizes and are centred on their middle pixel, which is the source position.

toycoronagraph.spatial.bandlimited_weights(nodes, x, band, n_taps=16)[source]#

Interpolation weights for a band-limited function, designed in the frequency domain.

Picks the n_taps nodes nearest to x and finds weights w (least squares) such that sum_j w_j exp(2 pi i f nodes_j) = exp(2 pi i f x) for every frequency |f| <= band. Any function whose spectrum lies within the band is then interpolated exactly (up to the fit residual), for uniform or non-uniform nodes. At a node, that node is returned with weight 1.

Parameters:

nodes (np.ndarray) – (n_nodes,) sample positions (any unit, e.g. pixels).
x (float) – Position to interpolate at, in the unit of nodes.
band (float) – Highest frequency present, in cycles per unit of the positions.
n_taps (int) – Number of nearest nodes used (all nodes if there are fewer).

Returns:

idx (np.ndarray): (m,) int indices into nodes, sorted (m = min(n_taps, n_nodes), or 1 when x is on a node).
w (np.ndarray): (m,) float64 weights; the value at x is w @ values[idx].

Return type:

tuple

toycoronagraph.spatial.positions_to_polar(img_pixel, positions)[source]#

Convert pixel positions to polar coordinates about the star.

The star sits at pixel coordinate c = (img_pixel-1)/2 on both axes (between the four central pixels for even img_pixel), the origin of the focal grid of the optics (optics.VortexBackend). The angle is measured from the +column axis towards the +row axis, the same frame as hcipy (x = column, y = row).

Parameters:

img_pixel (int) – Number of pixels on one edge of the image.
positions (np.ndarray) – (n, 2) array of (row, col) pixel positions (float or int).

Returns:

r (np.ndarray): (n,) float64 distance from the star, in pixels.
phi (np.ndarray): (n,) float64 angle, in radians in [-pi, pi].

Return type:

tuple

toycoronagraph.spatial.equal_error_grid(img_pixel, r_metric, d_r, d_phi, tol, max_spacing, n_theta_min=8, r_inner_max=None)[source]#

Build the PSF sample grid so neighbouring samples differ in PSF shape by about tol.

The grid follows how fast the (normalized) PSF shape actually changes: d_r(r) is the shape change per pixel of radius and d_phi(r) per radian of angle (see psf.shape_change_rates). Ring radii are placed at equal steps of the accumulated radial change (integral of d_r = tol between rings), up to the first ring beyond the image corner, and each ring gets ~2*pi*d_phi(r)/tol samples, so the shape difference between neighbours is ~tol everywhere. Two floors keep the far field well covered: ring spacing and the arc between samples are at most max_spacing pixels, and each ring has at least n_theta_min samples.

Near the star the required spacing drops below a pixel; inside that radius (r_pix) every pixel is its own sample, so those pixels get their exact PSF. Ring samples sit at their exact polar positions (not snapped to pixels): the convolution applies each kernel around the pixels it serves, so a sample’s position only sets the kernel’s shape. All samples of a ring then share one radius, and the ring structure stays regular for blending.

Every pixel gets up to four samples with blend weights (make_blend_weights); samples that no pixel uses are dropped.

Parameters:

img_pixel (int) – Number of pixels on one edge of the image.
r_metric (np.ndarray) – (m,) increasing radii, in pixels, at which the rates are given.
d_r (np.ndarray) – (m,) shape change per pixel of radius at r_metric.
d_phi (np.ndarray) – (m,) shape change per radian of angle at r_metric.
tol (float) – Target shape change (L1 of unit-sum kernels) between neighbouring samples.
max_spacing (float) – Largest ring spacing and arc between samples, in pixels.
n_theta_min (int) – Fewest samples on a ring.
r_inner_max (float) – Cap on r_pix, in pixels (bounds the number of per-pixel samples); None caps it at the image corner only.

Returns:

The grid, with n the number of samples kept and n_rings the number of rings:

positions ((n, 2) float64): sample positions (row, col), in pixels;
rings ((n,) int): ring label of each sample: 0 for a per-pixel inner sample, k+1 for ring k;
ring_radii ((n_rings,) float64): ring radii, in pixels (ring 0 at r_pix, the last beyond the image corner);
ring_theta0 ((n_rings,) float64): angle of the first sample of each ring, in radians;
ring_ntheta ((n_rings,) int): number of samples on each ring before pruning;
r_pix (float): radius, in pixels, inside which every pixel is its own sample;
inds ((img_pixel, img_pixel, 4) int), wts ((img_pixel, img_pixel, 4) float64): blend weights (see make_blend_weights), with inds re-indexed to the kept samples (0 where the weight is 0).
The ring_* arrays describe the grid before pruning (all rings, all samples).

Return type:

dict

toycoronagraph.spatial.make_blend_weights(img_pixel, grid)[source]#

Blend weights of the PSF samples for every pixel.

A pixel inside r_pix uses its own sample (weight 1). Any other pixel lies between two rings: on each ring the two samples bracketing the pixel’s angle are weighted linearly in angle, and the two rings linearly in radius, so up to four samples with weights summing to 1. A function that is linear in radius is reproduced exactly.

Parameters:

img_pixel (int) – Number of pixels on one edge of the image.
grid (dict) – Output of equal_error_grid before pruning: rings, ring_radii, ring_theta0, ring_ntheta and r_pix (samples listed as built there: inner pixels in row-major order, then ring after ring). Needs at least two rings if any pixel lies at or beyond r_pix.

Returns:

inds (np.ndarray): (img_pixel, img_pixel, 4) int sample indices; slots 0-1 on the inner ring, 2-3 on the outer ring (inner pixels: slot 0 only).
wts (np.ndarray): (img_pixel, img_pixel, 4) float64 weights, summing to 1 per pixel.

Return type:

tuple

toycoronagraph.spatial.blend_map(inds, wts, values)[source]#

Blend per-sample values onto the pixels, e.g. the coronagraph transmission map.

Parameters:

inds (np.ndarray) – (img_pixel, img_pixel, 4) int sample indices (from equal_error_grid).
wts (np.ndarray) – (img_pixel, img_pixel, 4) float blend weights (from equal_error_grid).
values (np.ndarray) – (n_samples,) one value per sample.

Returns:

(img_pixel, img_pixel) map, sum_k wts[…, k]*values[inds[…, k]].

Return type:

np.ndarray

toycoronagraph.spatial.make_psf_inds(img_pixel, positions)[source]#

Assign every pixel to its nearest PSF sample position (for convolve_with_spatial_psfs).

Parameters:

img_pixel (int) – Number of pixels on one edge of the image.
positions (np.ndarray) – (n, 2) array of (row, col) sample positions, in pixels.

Returns:

(img_pixel, img_pixel) int array of indices into positions.

Return type:

np.ndarray

toycoronagraph.spatial.psf_convolve_cpu(im, psf)[source]#

Convolve an image with a PSF on the CPU (scipy fftconvolve, mode ‘same’).

Parameters:

im (np.ndarray) – (ny, nx) image.
psf (np.ndarray) – (ky, kx) PSF kernel with odd dimensions, centred at its middle pixel.

Returns:

(ny, nx) convolved image.

Return type:

np.ndarray

toycoronagraph.spatial.psf_convolve_gpu(im, psf)[source]#

Convolve an image with a PSF on the GPU using CuPy (fftconvolve, mode ‘same’).

Parameters:

im (np.ndarray) – (ny, nx) image.
psf (np.ndarray) – (ky, kx) PSF kernel with odd dimensions, centred at its middle pixel.

Returns:

(ny, nx) convolved image, copied back to the host.

Return type:

np.ndarray

Raises:

ImportError – If CuPy is not installed.

toycoronagraph.spatial.convolve_with_spatial_psfs(im_in, psfs, psf_inds, patch_bounds=None, used_inds=None, coron_tmap=None, use_gpu=False, ncores=-2, method='auto')[source]#

Convolve an image where each pixel uses its own (nearest-sample) PSF.

The image is split by PSF index: the pixels assigned to PSF i are convolved with psfs[i] and the results are summed. ‘full’ convolves the whole frame once per PSF; ‘localized’ only the padded bounding box of each PSF’s pixels.

Parameters:

im_in (np.ndarray) – (ny, nx) image to be convolved (any real dtype; not modified).
psfs (np.ndarray) – (n_psf, ky, kx) spatially-sampled PSFs (odd ky, kx). Generally, each slice is the PSF of a source at one sampled position, recentred so that the source falls on the central pixel.
psf_inds (np.ndarray) – (ny, nx) int array; each entry is the index of the slice of psfs with which that pixel of im_in is convolved (e.g. make_psf_inds).
patch_bounds (np.ndarray) – (n_used, 4) pre-computed padded bounding boxes (from _get_psf_patch_bounds), matching used_inds. If it or used_inds is None, both are computed here when method is ‘auto’ or ‘localized’.
used_inds (np.ndarray) – (n_used,) pre-computed PSF indices present in psf_inds. None: computed here.
coron_tmap (np.ndarray) – (ny, nx) coronagraph transmission, by which im_in is multiplied before convolution; None for no transmission map.
use_gpu (bool) – Use CuPy for the FFT convolutions (then run serially).
ncores (int) – Number of threads (joblib n_jobs convention: -1 all, -2 all but one) when use_gpu is False; 1 runs serially.
method (str) – ‘auto’, ‘localized’ or ‘full’ (case-insensitive). ‘localized’ convolves padded bounding boxes around each PSF region, ‘full’ the whole frame for every PSF index, and ‘auto’ picks the cheaper by an FFT work estimate.

Returns:

(ny, nx) convolved image (float64 for real input).

Return type:

np.ndarray

Raises:

ValueError – If method is not one of the three.

toycoronagraph.spatial.convolve_with_blended_psfs(im_in, psfs, inds, wts, coron_tmap=None, use_gpu=False, ncores=-2)[source]#

Convolve an image where each pixel uses its own blend of spatially-sampled PSFs.

Pixel p contributes im[p] * coron_tmap[p] * sum_k wts[p, k] * psfs[inds[p, k]], centred on p. Samples used by a single pixel (the per-pixel samples near the star) are added directly; every other sample convolves the image weighted by its blend weights over the patch where they are non-zero. Kernels are never cropped by the image edge before being placed: the result is accumulated on a padded canvas and cut to the image at the end.

Parameters:

im_in (np.ndarray) – (ny, nx) image to be convolved (not modified).
psfs (np.ndarray) – (n_samples, K, K) sample kernels (odd K, centred; e.g. float32: each is convolved in float64).
inds (np.ndarray) – (ny, nx, 4) int sample indices of each pixel’s blend.
wts (np.ndarray) – (ny, nx, 4) float blend weights of each pixel.
coron_tmap (np.ndarray) – (ny, nx) coronagraph transmission, by which im_in is multiplied first; None for no transmission map.
use_gpu (bool) – Use CuPy for the FFT convolutions (then run serially).
ncores (int) – Number of threads (joblib n_jobs convention: -1 all, -2 all but one).

Returns:

(ny, nx) float64 convolved image.

Return type:

np.ndarray

toycoronagraph.spatial.spatial_iwa(model)[source]#

Inner working angle of spatially-sampled PSFs.

The IWA is the smallest source radius at which the core throughput (the PSF’s flux above half its peak, PSF = kernel x throughput) reaches half its maximum; the core throughput is made non-decreasing in radius first.

Parameters:: model (dict) – From psf.spatial_psf_calculation (uses img_pixel, positions, psfs and throughput).
Returns:: The IWA in pixels (comparable to psf.cir_iwa).
Return type:: float

toycoronagraph.spatial.spatial_psf(pre_img, model, iwa_ignore=False, use_gpu=False, ncores=-2, report=True)[source]#

Calculate the final image of any (not necessarily circularly symmetric) pre-image through spatially-varying coronagraphic PSFs.

Parameters:

pre_img (np.ndarray) – (img_pixel, img_pixel) input image (e.g. Jy/arcsec^2); the final image is in the same unit.
model (dict) – From psf.spatial_psf_calculation (psfs, inds, wts, tmap, img_pixel, …).
iwa_ignore (bool) – Set the target to 0 inside the IWA (spatial_iwa) first.
use_gpu (bool) – Use CuPy for the FFT convolutions.
ncores (int) – Number of threads (joblib n_jobs convention).
report (bool) – Print the estimated error of this image (see image_error_estimate), if the model has an error report.

Returns:

(img_pixel, img_pixel) float64 final image.

Return type:

np.ndarray

toycoronagraph.spatial.image_error_estimate(model, pre_img, verbose=True)[source]#

Estimated error of the final image of pre_img, from the build-time error check.

Every pixel’s light is spread by its effective PSF, whose error (end to end, relative to the PSF’s total flux) was measured per radius band when the PSFs were built (psf.spatial_error_report). Weighting those errors by the flux each pixel sends through the coronagraph (pre_img * tmap) gives the relative error of this image’s light: ‘typical’ from the checked pixels’ errors, interpolated in radius (or the band medians, for older models), ‘up to’ from the band maxima. The same weighting per error stage shows what dominates.

Parameters:

model (dict) – From psf.spatial_psf_calculation, with its error report (img_pixel, tmap, err_edges, err_stages, err_median, err_max; err_pix_r and err_pix if present; error_target, optional).
pre_img (np.ndarray) – (img_pixel, img_pixel) input image.
verbose (bool) – Print the estimate (and a warning above model[‘error_target’]).

Returns:

Relative errors (fractions, not %): typical (float), upper (float) and per_stage ({stage name: typical error}); empty if pre_img has no flux.

Return type:

dict