Appendix A — The helper module hnotools.py
The functions shared by all chapters. Everything here is something you should be able to write yourself.
"""Small helpers shared by the notebooks. Read this file: it is short, and
everything in it is something you should be able to write yourself."""
import os
import numpy as np
from numpy.fft import fft2, ifft2, ifftshift, fftfreq
import matplotlib.pyplot as plt
BLUE, ORANGE, AQUA = '#2a78d6', '#eb6834', '#1baf7a'
INK, INK2, SURFACE, GRIDC = '#0b0b0b', '#52514e', '#fcfcfb', '#e4e3df'
def style():
"""One plotting style for the whole series."""
plt.rcParams.update({
'figure.facecolor': SURFACE, 'axes.facecolor': SURFACE, 'savefig.facecolor': SURFACE,
'axes.edgecolor': GRIDC, 'axes.labelcolor': INK2, 'text.color': INK,
'xtick.color': INK2, 'ytick.color': INK2, 'axes.grid': True, 'grid.color': GRIDC,
'grid.linewidth': 0.6, 'axes.spines.top': False, 'axes.spines.right': False,
'font.size': 9, 'axes.titlesize': 9.5, 'lines.linewidth': 2.0,
'legend.frameon': False, 'figure.dpi': 110, 'image.cmap': 'gray',
'image.interpolation': 'nearest'})
def load_sample():
"""The real-data sample (see make_sample.py for what each array is)."""
here = os.path.dirname(os.path.abspath(__file__))
return dict(np.load(os.path.join(here, 'data', 'm16_sample.npz')))
def gaussian_psf(shape, fwhm):
"""Circular Gaussian PSF, centred on the array, summing to 1."""
sigma = fwhm / (2 * np.sqrt(2 * np.log(2)))
y, x = np.indices(shape)
y, x = y - shape[0] // 2, x - shape[1] // 2
p = np.exp(-(x**2 + y**2) / (2 * sigma**2))
return p / p.sum()
def moffat_psf(shape, fwhm_major, fwhm_minor=None, pa_deg=0.0, beta=2.5):
"""Elliptical Moffat PSF, the usual model of a star seen through the
atmosphere: (1 + r^2)^(-beta), heavier wings than a Gaussian."""
fwhm_minor = fwhm_minor or fwhm_major
c = 2 * np.sqrt(2 ** (1 / beta) - 1) # FWHM = c * alpha
a, b = fwhm_major / c, fwhm_minor / c
y, x = np.indices(shape)
y, x = y - shape[0] // 2, x - shape[1] // 2
t = np.radians(pa_deg)
u = (x * np.cos(t) + y * np.sin(t)) / a
v = (-x * np.sin(t) + y * np.cos(t)) / b
p = (1 + u**2 + v**2) ** (-beta)
return p / p.sum()
def psf2otf(psf):
"""Fourier transform of a centred PSF (its optical transfer function)."""
return fft2(ifftshift(psf))
def blur(x, psf):
"""Convolve x with psf, with periodic boundaries (one FFT product)."""
return np.real(ifft2(fft2(x) * psf2otf(psf)))
def radial_profile(a):
"""Azimuthal average of a 2-D Fourier-domain array (fft2 layout).
Returns frequency in cycles/pixel and the averaged values."""
n0, n1 = a.shape
r = np.hypot(fftfreq(n0)[:, None], fftfreq(n1)[None, :])
idx = np.rint(r * max(n0, n1)).astype(int)
nb = max(n0, n1) // 2 + 1
tot = np.bincount(idx.ravel(), weights=np.abs(a).ravel(), minlength=nb)[:nb]
cnt = np.bincount(idx.ravel(), minlength=nb)[:nb]
return np.arange(nb) / max(n0, n1), tot / np.maximum(cnt, 1)
def rel_error(estimate, truth):
"""Relative error ||estimate - truth|| / ||truth||."""
return np.linalg.norm(estimate - truth) / np.linalg.norm(truth)
def show(images, titles=None, lo=1, hi=99.7, same_scale=False, width=3.0, gamma=0.5):
"""Show images side by side with a percentile stretch.
same_scale=True uses the first image's limits for all of them."""
images = list(images)
fig, axes = plt.subplots(1, len(images), figsize=(width * len(images), width * 1.08), squeeze=False)
lim = None
for ax, im, k in zip(axes[0], images, range(len(images))):
if lim is None or not same_scale:
lim = np.percentile(im, [lo, hi])
z = np.clip((im - lim[0]) / (lim[1] - lim[0] + 1e-30), 0, 1) ** gamma
ax.imshow(z, vmin=0, vmax=1)
ax.set_xticks([]); ax.set_yticks([]); ax.grid(False)
for sp in ax.spines.values():
sp.set_visible(False)
if titles:
ax.set_title(titles[k])
fig.tight_layout(pad=0.3)
return fig
def find_stars(img, n=10, box=9):
"""Positions (row, col) and peak values of the n brightest local maxima,
brightest first. Good enough to locate stars in a small field."""
from scipy import ndimage
smooth = ndimage.gaussian_filter(img, 1.0)
peak = (smooth == ndimage.maximum_filter(smooth, box))
peak[:box, :] = peak[-box:, :] = False
peak[:, :box] = peak[:, -box:] = False
rows, cols = np.nonzero(peak)
order = np.argsort(img[rows, cols])[::-1][:n]
return [(int(rows[i]), int(cols[i]), float(img[rows[i], cols[i]])) for i in order]
def lecture_signal(N=250):
"""The 1-D test signal of H. Talbot's lecture notes (notebook 1):
steps plus slow waves, sampled on [0, 3 pi]. Returns (abscissa, signal)."""
s = np.linspace(0, 3 * np.pi, N)
return s, 1.0 * (np.cos(2 * s) > 0) + 2 * np.sin(s / 4) + 0.1 * np.cos(3 * s + 1)
def gaussian_kernel_1d(sigma, half):
"""Normalised 1-D Gaussian kernel of 2*half + 1 samples."""
o = np.arange(-half, half + 1)
k = np.exp(-o**2 / (2 * sigma**2))
return k / k.sum()
def blur_matrix_1d(kernel, N, mode='constant'):
"""Blurring matrix of a signal of length N (notebook 4): one column per
sample. mode = 'constant' (zero), 'wrap' (periodic) or 'reflect'."""
from scipy import ndimage
return np.column_stack([ndimage.convolve1d(np.eye(N)[j], kernel, mode=mode) for j in range(N)])
# ---- reflexive boundaries: the cosine transform (notebook 4) ----
def dct2(a):
from scipy.fft import dctn
return dctn(a, type=2, norm='ortho')
def idct2(a):
from scipy.fft import idctn
return idctn(a, type=2, norm='ortho')
def dct_eigenvalues(psf, shape):
"""Eigenvalues of the blur with reflexive boundaries, for a PSF that is
symmetric in both directions: dct2(A e1) / dct2(e1)."""
from scipy import ndimage
e1 = np.zeros(shape); e1[0, 0] = 1.0
return dct2(ndimage.convolve(e1, psf, mode='reflect')) / dct2(e1)
def make_test_problem(big, psf, noise, margin=64, seed=1):
"""Honest test data (notebook 4): blur a large scene, keep its centre,
add noise. Returns (x_true, b)."""
from scipy import ndimage
inner = slice(margin, big.shape[0] - margin)
blurred = ndimage.convolve(big, psf, mode='reflect')[inner, inner]
return big[inner, inner], blurred + noise * np.random.default_rng(seed).standard_normal(blurred.shape)
def taper(b, width=16):
"""Fade an image to its mean over `width` pixels along every edge, so
that periodic boundaries do no harm (notebook 4)."""
ry, rx = np.ones(b.shape[0]), np.ones(b.shape[1])
edge = 0.5 - 0.5 * np.cos(np.pi * np.arange(width) / width)
for r in (ry, rx):
r[:width] = edge; r[-width:] = edge[::-1]
return b.mean() + (b - b.mean()) * np.outer(ry, rx)
def fit_star(img, row, col, half=10):
"""Fit a circular Moffat profile on a constant background to the star
near (row, col) (notebook 3). Returns a dict with the FWHM in pixels."""
from scipy import optimize
stamp = img[row - half:row + half + 1, col - half:col + half + 1]
yy, xx = np.indices(stamp.shape) - half
def model(p):
amp, y0, x0, alpha, beta, bg = p
return bg + amp * (1 + ((yy - y0)**2 + (xx - x0)**2) / alpha**2) ** (-beta)
span = stamp.max() - np.median(stamp)
fit = optimize.least_squares(lambda p: (model(p) - stamp).ravel(),
x0=[span, 0, 0, 3.0, 2.5, np.median(stamp)],
bounds=([0, -3, -3, 0.3, 1.1, -np.inf], [10 * span + 1e-30, 3, 3, 30, 50, np.inf]))
amp, y0, x0, alpha, beta, bg = fit.x
return {'fwhm': 2 * alpha * np.sqrt(2 ** (1 / beta) - 1), 'beta': beta, 'amp': amp,
'bg': bg, 'row': row + y0, 'col': col + x0}