Source code for lenstronomy.LightModel.Profiles.shapelets_polar

__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'] 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} 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: roational invariance :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
def _chi_n_m(self, 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) def _index2n(self, 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] def poly2index(self, 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] def num_param(self, 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} 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: roational invariance :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
def _chi_n_m(self, 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./(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. * 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] def num_param(self, n_max): """ :param n_max: maximal polynomial order :return: number of basis components """ return int((n_max + 1)**2)
def _index2n(self, 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] def poly2index(self, 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} 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 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 x in range(n_max + 1)] for y 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 def _pre_calc_function(self, 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. * 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)