# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Classes for working with radial distributions,
e.g. the PSF or a source or a PSF-convolved source.
TODO: ThetaCalculator2D and ModelThetaCalculator are not
finished and need tests!
"""
from __future__ import absolute_import, division, print_function, unicode_literals
import numpy as np
__all__ = [
'ModelThetaCalculator',
'ThetaCalculator',
'ThetaCalculator2D',
'ThetaCalculatorScipy',
]
[docs]class ThetaCalculator(object):
"""Provides methods ``containment_fraction(theta)`` and ``containment_radius(containment_fraction)``
given some 1D distribution (not necessarily normalized).
Notes
If you have to compute theta or containment many times for
the same dist, this is much faster than ThetaCalculatorScipy.
If you want only one value it could actually be slower,
especially the containment calculation.
Parameters
----------
dist : callable
Distribution function dp / dtheta2 (theta2)
theta_max : float
Integration range will be 0 .. theta_max ^ 2
nbins : int
Integration step size
normalize : bool
Normalize discretized distribution to 1?
"""
def __init__(self, dist, theta_max, n_bins, normalize=False):
theta2 = np.linspace(0, theta_max ** 2, n_bins)
dtheta2 = theta2[1] - theta2[0]
p = (dist(theta2) * dtheta2).cumsum()
if normalize:
p /= p[-1]
self.theta2, self.p = theta2, p
[docs] def containment_fraction(self, theta):
"""Compute containment fraction for a given theta."""
theta = np.asarray(theta)
index = np.where(self.theta2 > theta ** 2)[0][0]
return self.p[index]
[docs] def containment_radius(self, containment_fraction):
"""Compute theta for a given containment fraction."""
containment_fraction = np.asarray(containment_fraction)
index = np.where(self.p > containment_fraction)[0][0]
return np.sqrt(self.theta2[index])
[docs]class ThetaCalculatorScipy(object):
"""Same functionality as NumericalThetaCalculator, but uses
``scipy.integrate.quad`` and ``scipy.optimize.fsolve`` instead.
Notes:
It is more precise than ThetaCalculator and doesn't
require you to think about which theta binning and range
gives your desired precision.
If you have to compute many thetas this can be quite slow
because it is a root finding with nested integration.
Parameters
----------
dist : callable
Probability distribution (probability per theta ^ 2)
theta_max : float
Integration range will be 0 .. theta_max ^ 2
normalize : bool
Normalize discretized distribution to 1?
"""
def __init__(self, dist, theta_max, normalize=False):
self.dist = dist
self.theta_max = theta_max
if normalize:
self.p_total = self.containment_fraction(theta_max, False)
else:
self.p_total = 1
[docs] def containment_fraction(self, theta):
from scipy.integrate import quad
p = quad(self.dist, 0, theta ** 2)[0]
return p / self.p_total
[docs] def containment_radius(self, containment_fraction):
"""Compute containment angle using the containment_fraction
method plus numerical root finding."""
from scipy.optimize import brentq
def f(theta):
return self.containment_fraction(theta) - containment_fraction
return brentq(f, 0, self.theta_max)
[docs]class ThetaCalculator2D(ThetaCalculatorScipy):
"""Methods to compute theta and containment
for a given 2D probability distribution image.
Typically this method is used for PSF-convolved
model images, where analytical distributions or
1D distributions are not available.
Note: The theta and containment is calculated relative
to the origin (x, y) = (0, 0).
Note: We do simple bin summing. In principle we could
do integration over bins by using scipy.integrate.dblquad
in combination with e.g. scipy.interpolate.interp2d,
but for the speed / precision we need this is overkill.
TODO: I just realized that probably the best thing to
do is to bin (x,y) -> theta2, make a spline interpolation
and then call ThetaCalculatorScipy!
Parameters
----------
dist : 2-dimensional array
Probability distribution (per dx * dy)
x : 2-dimensional array
Pixel ``x`` coordinate array. Must match shape of ``dist``.
x : 2-dimensional array
Pixel ``x`` coordinate array. Must match share of ``dist``.
"""
def __init__(self, dist, x, y):
self.dist = dist / dist.sum()
self.theta2 = x ** 2 + y ** 2
self.theta_max = np.sqrt(self.theta2.max())
[docs] def containment_fraction(self, theta):
mask = self.theta2 < theta ** 2
return self.dist[mask].sum()
[docs]class ModelThetaCalculator(ThetaCalculator):
"""Compute containment radius for given radially symmetric
source and psf as well as desired containment fraction.
Uses 2D images for the computation.
Slow but simple, so useful to check more complicated methods.
Source and PSF must be callable and return
dP/dtheta (TODO: or dP/dtheta^2?)
fov = field of view (deg)
binsz = bin size (deg)
The source is supposed to be contained in the FOV
even after PSF convolution.
"""
def __init__(self, source, psf, fov, binsz, call2d=False):
from scipy.ndimage import convolve
# Compute source and psf 2D images
y, x = np.mgrid[-fov:fov:binsz, -fov:fov:binsz]
theta2 = x * x + y * y
if call2d:
source_image = source(x, y)
psf_image = psf(x, y)
else:
source_image = source.dpdtheta2(theta2)
psf_image = psf.dpdtheta2(theta2)
# Compute convolved image and normalize it
p = convolve(source_image, psf_image)
p /= p.sum()
# Store the theta2 and p arrays
self.p, self.theta2 = p, theta2