__author__ = "sibirrer"
import numpy as np
import math
import scipy.special
import lenstronomy.Util.param_util as param_util
from lenstronomy.Util.package_util import exporter
export, __all__ = exporter()
[docs]
@export
class ShapeletsPolar(object):
"""2D polar Shapelets, see Massey & Refregier 2005."""
param_names = ["amp", "beta", "n", "m", "center_x", "center_y"]
param_names_latex = {r"$I_0$", r"$\beta$", r"$n$", r"$m$", r"$x_0$", r"$y_0$"}
lower_limit_default = {
"amp": 0,
"beta": 0,
"n": 0,
"m": 0,
"center_x": -100,
"center_y": -100,
}
upper_limit_default = {
"amp": 100,
"beta": 100,
"n": 150,
"m": 150,
"center_x": 100,
"center_y": 100,
}
[docs]
def __init__(self):
"""Load interpolation of the Hermite polynomials in a range [-30,30] in order
n<= 150 :return:"""
pass
[docs]
def function(self, x, y, amp, beta, n, m, complex_bool, center_x, center_y):
"""
:param x: x-coordinate, numpy array
:param y: y-ccordinate, numpy array
:param amp: amplitude normalization
:param beta: shaplet scale
:param n: order of polynomial
:param m: rotational invariance
:param complex_bool: boolean; if True uses complex value of function _chi_n_m()
:param center_x: center of shapelet
:param center_y: center of shapelet
:return: amplitude of shapelet at possition (x, y)
"""
r, phi = param_util.cart2polar(x, y, center_x, center_y)
if complex_bool is True:
return amp * self._chi_n_m(r, beta, n, m) * np.exp(-1j * m * phi).imag
else:
return amp * self._chi_n_m(r, beta, n, m) * np.exp(-1j * m * phi).real
@staticmethod
def _chi_n_m(r, beta, n, m):
"""
:param r: radius / beta
:param beta: shapelet scale
:param n: non-negative integer
:param m: integer, running from -n to n in steps of two
:return: value of function (8) in Massey & Refregier, complex numbers
"""
m_abs = int(abs(m))
p = int((n - m_abs) / 2)
p2 = int((n + m_abs) / 2)
if p % 2 == 0: # if p is even
prefac = 1
else:
prefac = -1
prefactor = (
prefac
/ beta ** (m_abs + 1)
* np.sqrt(math.factorial(p) / (np.pi * math.factorial(p2)))
)
poly = scipy.special.genlaguerre(
n=p, alpha=m_abs
) # lower part, upper part of definition in Massey & Refregier
r_ = (r / beta) ** 2
l_n_alpha = poly(r_)
return prefactor * r**m_abs * l_n_alpha * np.exp(-((r / beta) ** 2) / 2)
@staticmethod
def _index2n(index):
"""
:param index: index of convention
:return: n order of Laguerre
"""
n_float = (-3 + np.sqrt(9 + 8 * index)) / 2
n_int = int(n_float)
if n_int == n_float:
n = n_int
else:
n = n_int + 1
return n
[docs]
def index2poly(self, index):
"""Manages the convention from an iterative index to the specific polynomial n,
m, (real/imaginary part)
:param index: int, index of list
:return: n, m bool
"""
n = self._index2n(index)
num_prev = n * (n + 1) / 2
num = index + 1
delta = int(num - num_prev - 1)
if n % 2 == 0:
if delta == 0:
m = delta
complex_bool = False
elif delta % 2 == 0:
complex_bool = True
m = delta
else:
complex_bool = False
m = delta + 1
else:
if delta % 2 == 0:
complex_bool = False
m = delta + 1
else:
complex_bool = True
m = delta
return n, m, complex_bool
[docs]
@staticmethod
def poly2index(n, m, complex_bool):
"""
:param n: non-negative integer
:param m: integer, running from -n to n in steps of two
:param complex_bool: bool, if True, assigns complex part
:return:
"""
index = n * (n + 1) / 2
if complex_bool is True:
if m == 0:
raise ValueError("m=0 can not have imaginary part!")
if n % 2 == 0:
if m % 2 == 0:
if m == 0:
index += 1
else:
index += m
if complex_bool is True:
index += 1
else:
raise ValueError("m needs to be even for even n!")
else:
if complex_bool is True:
index += m + 1
else:
index += m
return int(index - 1)
[docs]
@staticmethod
def num_param(n_max):
"""
:param n_max: maximal polynomial order
:return: number of basis components
"""
return int((n_max + 1) * (n_max + 2) / 2)
[docs]
@export
class ShapeletsPolarExp(object):
"""2D exponential shapelets, Berge et al.
2019
"""
param_names = ["amp", "beta", "n", "m", "center_x", "center_y"]
lower_limit_default = {
"amp": 0,
"beta": 0,
"n": 0,
"m": 0,
"center_x": -100,
"center_y": -100,
}
upper_limit_default = {
"amp": 100,
"beta": 100,
"n": 150,
"m": 150,
"center_x": 100,
"center_y": 100,
}
[docs]
def __init__(self):
"""Load interpolation of the Hermite polynomials in a range [-30,30] in order
n<= 150 :return:"""
pass
[docs]
def function(self, x, y, amp, beta, n, m, complex_bool, center_x, center_y):
"""
:param x: x-coordinate, numpy array
:param y: y-ccordinate, numpy array
:param amp: amplitude normalization
:param beta: shaplet scale
:param n: order of polynomial
:param m: rotational invariance
:param complex_bool: boolean; if True uses complex value of function _chi_n_m()
:param center_x: center of shapelet
:param center_y: center of shapelet
:return: amplitude of shapelet at possition (x, y)
"""
r, phi = param_util.cart2polar(x, y, center_x, center_y)
if complex_bool is True:
return amp * self._chi_n_m(r, beta, n, m) * np.exp(-1j * m * phi).imag
else:
return amp * self._chi_n_m(r, beta, n, m) * np.exp(-1j * m * phi).real
@staticmethod
def _chi_n_m(r, beta, n, m):
"""
:param r: radius
:param beta: shapelet scale
:param n: non-negative integer
:param m: integer, running from -n to n in steps of two
:return: value of function (20) in Berge et al. 2019 without complex coefficient
"""
m_abs = int(abs(m))
p = int(n - m_abs)
p2 = int(n + m_abs)
if p % 2 == 0: # if p is even
prefac = 1
else:
prefac = -1
prefactor = prefac * np.sqrt(
2.0
/ (beta * np.pi * (2 * n + 1) ** 3)
* math.factorial(p)
/ math.factorial(p2)
)
poly = scipy.special.genlaguerre(
n=p, alpha=2 * m_abs
) # lower part, upper part of definition in Massey & Refregier
x_ = 2.0 * r / (beta * (2 * n + 1))
l_n_alpha = poly(x_)
return prefactor * x_**m_abs * l_n_alpha * np.exp(-x_ / 2) / np.sqrt(beta)
[docs]
@staticmethod
def num_param(n_max):
"""
:param n_max: maximal polynomial order
:return: number of basis components
"""
return int((n_max + 1) ** 2)
@staticmethod
def _index2n(index):
"""
:param index: index convention
:return: n
"""
n_float = np.sqrt(index + 1) - 1
n_int = int(n_float)
if n_int == n_float:
n = n_int
else:
n = n_int + 1
return n
[docs]
def index2poly(self, index):
"""
:param index:
:return:
"""
n = self._index2n(index)
num_prev = n**2
delta = index + 1 - num_prev
if delta % 2 == 1:
m = int((delta - 1) / 2)
complex_bool = False
else:
m = int(delta / 2)
complex_bool = True
return n, m, complex_bool
[docs]
@staticmethod
def poly2index(n, m, complex_bool):
"""
:param n:
:param m:
:param complex_bool:
:return: index convention, integer
"""
index = n**2
index += 2 * m
if complex_bool is True:
index -= 1
return int(index)
[docs]
@export
class ShapeletSetPolar(object):
"""Class to operate on entire shapelet set."""
param_names = ["amp", "n_max", "beta", "center_x", "center_y"]
lower_limit_default = {"beta": 0, "center_x": -100, "center_y": -100}
upper_limit_default = {"beta": 100, "center_x": 100, "center_y": 100}
[docs]
def __init__(self, exponential=False):
if exponential is True:
self.shapelets = ShapeletsPolarExp()
else:
self.shapelets = ShapeletsPolar()
self._exponential = exponential
[docs]
def function(self, x, y, amp, n_max, beta, center_x=0, center_y=0):
"""
:param x:
:param y:
:param amp:
:param n_max:
:param beta:
:param center_x:
:param center_y:
:return:
"""
num_param = self.shapelets.num_param(n_max)
f_ = np.zeros(len(np.atleast_1d(x)))
l_list = self._pre_calc(x, y, beta, n_max, center_x, center_y)
for i in range(num_param):
out = self._pre_calc_function(l_list, i) * amp[i]
f_ += out
try:
len(x)
except:
f_ = f_[0]
return np.nan_to_num(f_)
[docs]
def function_split(self, x, y, amp, n_max, beta, center_x=0, center_y=0):
num_param = self.shapelets.num_param(n_max)
A = []
L_list = self._pre_calc(x, y, beta, n_max, center_x, center_y)
for i in range(num_param):
A.append(self._pre_calc_function(L_list, i) * amp[i])
return A
def _pre_calc(self, x, y, beta, n_max, center_x, center_y):
"""
:param x:
:param y:
:param beta:
:param n_max:
:param center_x:
:param center_y:
:return:
"""
L_list = []
# polar coordinates
r, phi = param_util.cart2polar(x, y, center_x, center_y)
num_param = self.shapelets.num_param(n_max)
# compute real and imaginary part of the complex angles in the range n_max
theta_m_real_list, theta_m_imag_list = [], []
for i in range(n_max + 1):
m = i
exp_complex = np.exp(-1j * m * phi)
theta_m_real_list.append(exp_complex.real)
theta_m_imag_list.append(exp_complex.imag)
# compute the Laguerre polynomials in n, m
chi_n_m_list = [[0 for _ in range(n_max + 1)] for _ in range(n_max + 1)]
for n in range(n_max + 1):
for m in range(n + 1):
if (n - m) % 2 == 0 or self._exponential is True:
chi_n_m_list[n][m] = self.shapelets._chi_n_m(r, beta, n, m)
# combine together the pre-computed components
for index in range(num_param):
n, m, complex_bool = self.shapelets.index2poly(index)
if complex_bool is True:
L_i = chi_n_m_list[n][m] * theta_m_imag_list[m]
else:
L_i = chi_n_m_list[n][m] * theta_m_real_list[m]
L_list.append(L_i)
return L_list
@staticmethod
def _pre_calc_function(L_list, i):
"""Evaluates the shapelet function based on the pre-calculated components in
_pre_calc()
:param L_list: pre-calculated components
:param i: index conventions of the sequence of basis components
:return: shaplet basis at the pre-calculated positions
"""
return L_list[i]
[docs]
def decomposition(self, image, x, y, n_max, beta, deltaPix, center_x=0, center_y=0):
"""Decomposes an image into the shapelet coefficients in same order as for the
function call :param image:
:param x:
:param y:
:param n_max:
:param beta:
:param center_x:
:param center_y:
:return:
"""
num_param = self.shapelets.num_param(n_max)
param_list = np.zeros(num_param)
amp_norm = 1.0 * deltaPix**2
L_list = self._pre_calc(x, y, beta, n_max, center_x, center_y)
for i in range(num_param):
base = self._pre_calc_function(L_list, i) * amp_norm
param = np.sum(image * base)
n, m, complex_bool = self.shapelets.index2poly(i)
if m != 0:
param *= 2
param_list[i] = param
return param_list
[docs]
def index2poly(self, index):
"""
:param index: index of coefficient in the convention here
:return: n, m, complex_bool
"""
return self.shapelets.index2poly(index)