#!/usr/bin/env python
"""
subpixel.py -- localize PSFs to subpixel resolution
"""
# Numeric
import numpy as np
# Dataframes
import pandas as pd
# Image processing
from scipy import ndimage as ndi
# Low-level subpixel localization utilities
from .helper import (
assign_methods,
ring_mean,
rs,
I0_is_crazy,
estimate_I0,
estimate_I0_multiple_points,
estimate_snr,
fit_ls_int_gaussian,
fit_ls_point_gaussian,
fit_poisson_int_gaussian,
check_2d_gauss_fit,
)
[docs]def centroid(I, sub_bg=False):
"""
Find the center of a spot by taking its center
of mass.
args
----
I : 2D ndarray (YX), spot image
sub_bg : bool, subtract background prior to
taking center of mass
returns
-------
dict {
y : y centroid (pixels),
x : x centroid (pixels),
bg : estimated background intensity per
pixel (AU),
I0 : integrated spot intensity above
background (AU),
error_flag : int, the error code. 0 indicates
no errors.
}
"""
# Estimate background
bg = ring_mean(I)
# Subtract background
I_sub = np.clip(I - bg, 0.0, np.inf)
# Integrated intensity above background
I0 = I_sub.sum()
# Find spot centers
if sub_bg:
y, x = ndi.center_of_mass(I_sub)
else:
y, x = ndi.center_of_mass(I)
# Estimate SNR
snr = estimate_snr(I, I0)
# Return parameter estimates
return dict(
(
("y", y),
("x", x),
("I0", I0),
("bg", bg),
("error_flag", 0),
("snr", snr),
)
)
[docs]def radial_symmetry(I, sigma=1.0, **kwargs):
"""
Estimate the center of a spot by the radial symmetry
method, described in Parthasarathy et al. Nat Met 2012.
Also infer the intensity of the spots assuming an
integrated Gaussian PSF. This is useful as a first guess
for iterative localization techniques.
args
----
I : 2D ndarray, spot image
sigma : float, static integrated 2D Gaussian
width
returns
-------
dict {
y : estimated y center (pixels),
x : estimated x center (pixels),
I0 : estimated PSF intensity (AU),
bg : estimated background intensity per
pixel (AU),
error_flag : int, the error flag. 0
if there are no errors,
}
"""
# Estimate spot centers
y, x = rs(I)
# Estimate background level
bg = ring_mean(I)
# Estimate Gaussian intensity
I0 = estimate_I0_multiple_points(I, y, x, bg, sigma=sigma)
# If I0 is crazy, use a different guess
if I0_is_crazy(I0):
I0 = np.clip(I - bg, 0.0, np.inf).sum()
error_flag = 1
else:
error_flag = 0
# Estimate SNR
snr = estimate_snr(I, I0)
# Return parameter estimate
return dict(
(
("y", y),
("x", x),
("I0", I0),
("bg", bg),
("error_flag", error_flag),
("snr", snr),
)
)
[docs]def ls_int_gaussian(
I,
sigma=1.0,
ridge=0.0001,
max_iter=10,
damp=0.3,
convergence=1.0e-4,
divergence=1.0,
):
"""
Estimate the maximum likelihood parameters for a symmetric
2D integrated Gaussian PSF model, given an observed spot
*I* with normally-distributed noise.
This method uses radial symmetry for a first guess, followed
by a Levenberg-Marquardt routine for refinement. The core
calculation is performed in quot.helper.fit_ls_int_gaussian.
args
----
I : 2D ndarray (YX), the observed spot
sigma : float, static Gaussian width
ridge : float, initial regularization term
for inversion of the Hessian
max_iter : int, the maximum tolerated number
of iterations
damp : float, damping factor for the update
at each iteration. Larger means faster
convergence, but less stable.
convergence : float, maximum magnitude of the update
vector at which to call convergence.
divergence : float, divergence criterion
returns
-------
dict, the parameter estimate, error estimates, and
related parameters about the fitting problem
(see below). Some of these can be useful for
QC.
"""
# Make the initial guess by radial symmetry
guess = np.array([*rs(I), 0.0, ring_mean(I)])
guess[2] = estimate_I0(I, guess[0], guess[1], guess[3], sigma=sigma)
# Run the fitting routine
pars, err, H_det, rmse, n_iter = fit_ls_int_gaussian(
I,
guess,
sigma=sigma,
ridge=ridge,
max_iter=max_iter,
damp=damp,
convergence=convergence,
divergence=divergence,
)
# Check for crazy fits
if check_2d_gauss_fit(I.shape, pars):
error_flag = 0
else:
error_flag = 1
# Estimate SNR
snr = estimate_snr(I, pars[2])
# Return parameter estimate
return dict(
(
("y", pars[0]),
("x", pars[1]),
("I0", pars[2]),
("bg", pars[3]),
("y_err", err[0]),
("x_err", err[1]),
("I0_err", err[2]),
("bg_err", err[3]),
("H_det", H_det),
("error_flag", error_flag),
("snr", snr),
("rmse", rmse),
("n_iter", n_iter),
)
)
[docs]def ls_point_gaussian(
I,
sigma=1.0,
ridge=0.0001,
max_iter=10,
damp=0.3,
convergence=1.0e-4,
divergence=1.0,
):
"""
Estimate the maximum likelihood parameters for a symmetric
2D pointwise-evaluated Gaussian PSF model, given an observed
spot *I* with normally-distributed noise.
Both integrated and pointwise-evaluated models have the same
underlying model: a symmetric, 2D Gaussian PSF. The
distinction is how they handle sampling on discrete pixels:
- the point Gaussian takes the value on each pixel to
be equal to the PSF function evaluated at the center
of that pixel
- the integrated Gaussian takes the value on each pixel
to be equal to the PSF integrated across the whole
area of the pixel
Integrated Gaussian models are more accurate and less prone
to edge biases, at the cost of increased complexity and
perhaps lower speed.
This method uses radial symmetry for a first guess, followed
by a Levenberg-Marquardt routine for refinement. The core
calculation is performed in quot.helper.fit_ls_point_gaussian.
args
----
I : 2D ndarray (YX), the observed spot
sigma : float, static Gaussian width
ridge : float, initial regularization term
for inversion of the Hessian
max_iter : int, the maximum tolerated number
of iterations
damp : float, damping factor for the update
at each iteration. Larger means faster
convergence, but less stable.
convergence : float, maximum magnitude of the update
vector at which to call convergence.
divergence : float, divergence criterion
returns
-------
dict, the parameter estimate, error estimates, and
related parameters about the fitting problem
(see below). Some of these can be useful for
QC.
"""
# Make the initial guess by radial symmetry
guess = np.array([*rs(I), 0.0, ring_mean(I)])
guess[2] = estimate_I0(I, guess[0], guess[1], guess[3], sigma=sigma)
# Run the fitting routine
pars, err, H_det, rmse, n_iter = fit_ls_point_gaussian(
I,
guess,
sigma=sigma,
ridge=ridge,
max_iter=max_iter,
damp=damp,
convergence=convergence,
divergence=divergence,
)
# Check for crazy fits
if check_2d_gauss_fit(I.shape, pars):
error_flag = 0
else:
error_flag = 1
# Estimate SNR
snr = estimate_snr(I, pars[2])
# Return parameter estimate
return dict(
(
("y", pars[0]),
("x", pars[1]),
("I0", pars[2]),
("bg", pars[3]),
("y_err", err[0]),
("x_err", err[1]),
("I0_err", err[2]),
("bg_err", err[3]),
("H_det", H_det),
("error_flag", error_flag),
("snr", snr),
("rmse", rmse),
("n_iter", n_iter),
)
)
[docs]def poisson_int_gaussian(
I,
sigma=1.0,
ridge=0.0001,
max_iter=10,
damp=0.3,
convergence=1.0e-4,
divergence=1.0,
):
"""
Estimate the maximum likelihood parameters for a symmetric
2D integrated Gaussian PSF model, given an observed spot
*I* with Poisson-distributed noise.
This method uses radial symmetry for a first guess, followed
by a Levenberg-Marquardt routine for refinement. The core
calculation is performed in quot.helper.fit_poisson_int_gaussian.
args
----
I : 2D ndarray (YX), the observed spot
sigma : float, static Gaussian width
ridge : float, initial regularization term
for inversion of the Hessian
max_iter : int, the maximum tolerated number
of iterations
damp : float, damping factor for the update
at each iteration. Larger means faster
convergence, but less stable.
convergence : float, maximum magnitude of the update
vector at which to call convergence.
divergence : float, divergence criterion
returns
-------
dict, the parameter estimate, error estimates, and
related parameters about the fitting problem
(see below). Some of these can be useful for
QC.
"""
# Make the initial guess by radial symmetry
guess = np.array([*rs(I), 0.0, ring_mean(I)])
guess[2] = estimate_I0(I, guess[0], guess[1], guess[3], sigma=sigma)
# Run the fitting routine
pars, err, H_det, rmse, n_iter = fit_poisson_int_gaussian(
I,
guess,
sigma=sigma,
ridge=ridge,
max_iter=max_iter,
damp=damp,
convergence=convergence,
divergence=divergence,
)
# Check for crazy fits
if check_2d_gauss_fit(I.shape, pars):
error_flag = 0
else:
error_flag = 1
# Estimate SNR
snr = estimate_snr(I, pars[2])
# Return parameter estimate
return dict(
(
("y", pars[0]),
("x", pars[1]),
("I0", pars[2]),
("bg", pars[3]),
("y_err", err[0]),
("x_err", err[1]),
("I0_err", err[2]),
("bg_err", err[3]),
("H_det", H_det),
("error_flag", error_flag),
("snr", snr),
("rmse", rmse),
("n_iter", n_iter),
)
)
##########################################
## MAIN SUBPIXEL LOCALIZATION FUNCTIONS ##
##########################################
# All localization methods available
METHODS = {
"centroid": centroid,
"radial_symmetry": radial_symmetry,
"ls_int_gaussian": ls_int_gaussian,
"ls_point_gaussian": ls_point_gaussian,
"poisson_int_gaussian": poisson_int_gaussian,
}
# Wrapper for all localization methods to run on
# a single PSF.
# Example usage:
# fit_pars = localize(psf_img, method='ls_int_gaussian',
# **method_kwargs)
localize = assign_methods(METHODS)(lambda I, **kwargs: None)
# Wrapper for all localization methods on a frame
# with several detections.
[docs]def localize_frame(
img,
positions,
method=None,
window_size=9,
camera_bg=0.0,
camera_gain=1.0,
**method_kwargs
):
"""
Run localization on multiple spots in a large 2D image,
returning the result as a pandas DataFrame.
args
----
img : 2D ndarray (YX), the image frame
positions : 2D ndarray of shape (n_spots, 2),
the y and x positions at which to
localize spots
method : str, a method in METHODS
window_size : int, the fitting window size
camera_bg : float, the BG per subpixel in the camera
camera_gain : float, the camera gain (grayvalues/photon)
method_kwargs: to the localization method
returns
-------
pandas.DataFrame with columns ['y_detect', 'x_detect']
plus all of the outputs from the localization function
"""
# If no method passed, return the unmodified dataframe
if method is None:
return positions
# Remove detections too close to the edge for a square
# subwindow
if len(positions.shape) == 2:
hw = window_size // 2
positions = positions[
(positions[:, 0] >= hw)
& (positions[:, 0] < img.shape[0] - hw)
& (positions[:, 1] >= hw)
& (positions[:, 1] < img.shape[1] - hw),
:,
]
# Get the localization method
method_f = METHODS.get(method)
# Localize a PSF in a subwindow of the image
def localize_subwindow(yd, xd):
psf_img = (
np.clip(
img[yd - hw : yd + hw + 1, xd - hw : xd + hw + 1]
- camera_bg,
0,
np.inf,
)
/ camera_gain
)
r = method_f(psf_img, **method_kwargs)
r.update({"y_detect": yd, "x_detect": xd})
r["y"] = r["y"] + yd - hw
r["x"] = r["x"] + xd - hw
return r
# Run localization on all PSF subwindows
result = pd.DataFrame(
[localize_subwindow(yd, xd) for yd, xd in positions]
)
return result
else:
return pd.DataFrame([])