lenstronomy.LensModel.Profiles.shapelet_pot_cartesian — lenstronomy 1.9.1 documentation

# Source code for lenstronomy.LensModel.Profiles.shapelet_pot_cartesian

    __author__ = 'sibirrer'

    # description of the polar shapelets in potential space

    import numpy as np
    import math
    import numpy.polynomial.hermite as hermite
    from lenstronomy.LensModel.Profiles.base_profile import LensProfileBase

    __all__ = ['CartShapelets']

    class CartShapelets(LensProfileBase):
        """
        this class contains the function and the derivatives of the cartesian shapelets
        """
        param_names = ['coeffs', 'beta', 'center_x', 'center_y']
        lower_limit_default = {'coeffs': [0], 'beta': 0, 'center_x': -100, 'center_y': -100}
        upper_limit_default = {'coeffs': [100], 'beta': 100, 'center_x': 100, 'center_y': 100}

        def function(self, x, y, coeffs, beta, center_x=0, center_y=0):
            shapelets = self._createShapelet(coeffs)
            n_order = self._get_num_n(len(coeffs))
            n = len(np.atleast_1d(x))
            if n <= 1:
                f_ = self._shapeletOutput(x, y, beta, shapelets, precalc=False)
            else:
                H_x, H_y = self.pre_calc(x, y, beta, n_order, center_x, center_y)
                f_ = self._shapeletOutput(H_x, H_y, beta, shapelets)
            return f_

        def derivatives(self, x, y, coeffs, beta, center_x=0, center_y=0):
            """
            returns df/dx and df/dy of the function
            """
            shapelets = self._createShapelet(coeffs)
            n_order = self._get_num_n(len(coeffs))
            dx_shapelets = self._dx_shapelets(shapelets, beta)
            dy_shapelets = self._dy_shapelets(shapelets, beta)
            n = len(np.atleast_1d(x))
            if n <= 1:
                f_x = self._shapeletOutput(x, y, beta, dx_shapelets, precalc=False)
                f_y = self._shapeletOutput(x, y, beta, dy_shapelets, precalc=False)
            else:
                H_x, H_y = self.pre_calc(x, y, beta, n_order+1, center_x, center_y)
                f_x = self._shapeletOutput(H_x, H_y, beta, dx_shapelets)
                f_y = self._shapeletOutput(H_x, H_y, beta, dy_shapelets)
            return f_x, f_y

        def hessian(self, x, y, coeffs, beta, center_x=0, center_y=0):
            """
            returns Hessian matrix of function d^2f/dx^2, d^2/dxdy, d^2/dydx, d^f/dy^2
            """
            shapelets = self._createShapelet(coeffs)
            n_order = self._get_num_n(len(coeffs))
            dxx_shapelets = self._dxx_shapelets(shapelets, beta)
            dyy_shapelets = self._dyy_shapelets(shapelets, beta)
            dxy_shapelets = self._dxy_shapelets(shapelets, beta)
            n = len(np.atleast_1d(x))
            if n <= 1:
                f_xx = self._shapeletOutput(x, y, beta, dxx_shapelets, precalc=False)
                f_yy = self._shapeletOutput(x, y, beta, dyy_shapelets, precalc=False)
                f_xy = self._shapeletOutput(x, y, beta, dxy_shapelets, precalc=False)
            else:
                H_x, H_y = self.pre_calc(x, y, beta, n_order+2, center_x, center_y)
                f_xx = self._shapeletOutput(H_x, H_y, beta, dxx_shapelets)
                f_yy = self._shapeletOutput(H_x, H_y, beta, dyy_shapelets)
                f_xy = self._shapeletOutput(H_x, H_y, beta, dxy_shapelets)
            return f_xx, f_xy, f_xy, f_yy

        def _createShapelet(self, coeffs):
            """
            returns a shapelet array out of the coefficients *a, up to order l

            :param num_l: order of shapelets
            :type num_l: int.
            :param coeff: shapelet coefficients
            :type coeff: floats
            :returns:  complex array
            :raises: AttributeError, KeyError
            """
            n_coeffs = len(coeffs)
            num_n = self._get_num_n(n_coeffs)
            shapelets=np.zeros((num_n+1, num_n+1))
            n = 0
            k = 0
            for coeff in coeffs:
                shapelets[n-k][k] = coeff
                k += 1
                if k == n + 1:
                    n += 1
                    k = 0
            return shapelets

        def _shapeletOutput(self, x, y, beta, shapelets, precalc=True):
            """
            returns the the numerical values of a set of shapelets at polar coordinates
            :param shapelets: set of shapelets [l=,r=,a_lr=]
            :type shapelets: array of size (n,3)
            :param coordPolar: set of coordinates in polar units
            :type coordPolar: array of size (n,2)
            :returns:  array of same size with coords [r,phi]
            :raises: AttributeError, KeyError
            """
            n = len(np.atleast_1d(x))
            if n <= 1:
                values = 0.
            else:
                values = np.zeros(len(x[0]))
            n = 0
            k = 0
            i = 0
            num_n = len(shapelets)
            while i < num_n * (num_n+1)/2:
                values += self._function(x, y, shapelets[n-k][k], beta, n-k, k, precalc=precalc)
                k += 1
                if k == n + 1:
                    n += 1
                    k = 0
                i += 1
            return values

        def _function(self, x, y, amp, beta, n1, n2, center_x=0, center_y=0, precalc=False):
            """

            :param amp: amplitude of shapelet
            :param beta: scale factor of shapelet
            :param n1: x-order
            :param n2: y-order
            :param center_x: center in x
            :param center_y: center in y
            :return:
            """
            if precalc:
                return amp * x[n1] * y[n2] / beta
            x_ = x - center_x
            y_ = y - center_y
            return amp * self.phi_n(n1, x_/beta) * self.phi_n(n2, y_/beta) /beta

        def _dx_shapelets(self, shapelets, beta):
            """
            computes the derivative d/dx of the shapelet coeffs
            :param shapelets:
            :param beta:
            :return:
            """
            num_n = len(shapelets)
            dx = np.zeros((num_n+1, num_n+1))
            for n1 in range(num_n):
                for n2 in range(num_n):
                    amp = shapelets[n1][n2]
                    dx[n1+1][n2] -= np.sqrt((n1+1)/2.) * amp
                    if n1 > 0:
                        dx[n1-1][n2] += np.sqrt(n1/2.) * amp
            return dx/beta

        def _dy_shapelets(self, shapelets, beta):
            """
            computes the derivative d/dx of the shapelet coeffs
            :param shapelets:
            :param beta:
            :return:
            """
            num_n = len(shapelets)
            dy = np.zeros((num_n+1, num_n+1))
            for n1 in range(num_n):
                for n2 in range(num_n):
                    amp = shapelets[n1][n2]
                    dy[n1][n2+1] -= np.sqrt((n2+1)/2.) * amp
                    if n2 > 0:
                        dy[n1][n2-1] += np.sqrt(n2/2.) * amp
            return dy/beta

        def _dxx_shapelets(self, shapelets, beta):
            dx_shapelets = self._dx_shapelets(shapelets, beta)
            return self._dx_shapelets(dx_shapelets, beta)

        def _dyy_shapelets(self, shapelets, beta):
            dy_shapelets = self._dy_shapelets(shapelets, beta)
            return self._dy_shapelets(dy_shapelets, beta)

        def _dxy_shapelets(self, shapelets, beta):
            dy_shapelets = self._dy_shapelets(shapelets, beta)
            return self._dx_shapelets(dy_shapelets, beta)

        def H_n(self, n, x):
            """
            constructs the Hermite polynomial of order n at position x (dimensionless)

            :param n: The n'the basis function.
            :type name: int.
            :param x: 1-dim position (dimensionless)
            :type state: float or numpy array.
            :returns:  array-- H_n(x).
            :raises: AttributeError, KeyError
            """
            n_array = np.zeros(n+1)
            n_array[n] = 1
            return hermite.hermval(x, n_array, tensor=False) #attention, this routine calculates every single hermite polynomial and multiplies it with zero (exept the right one)

        def phi_n(self,n,x):
            """
            constructs the 1-dim basis function (formula (1) in Refregier et al. 2001)

            :param n: The n'the basis function.
            :type name: int.
            :param x: 1-dim position (dimensionless)
            :type state: float or numpy array.
            :returns:  array-- phi_n(x).
            :raises: AttributeError, KeyError
            """
            prefactor = 1./np.sqrt(2**n*np.sqrt(np.pi)*math.factorial(n))
            return prefactor*self.H_n(n,x)*np.exp(-x**2/2.)

        def pre_calc(self, x, y, beta, n_order, center_x, center_y):
            """
            calculates the H_n(x) and H_n(y) for a given x-array and y-array
            :param x:
            :param y:
            :param amp:
            :param beta:
            :param n_order:
            :param center_x:
            :param center_y:
            :return: list of H_n(x) and H_n(y)
            """

            n = len(np.atleast_1d(x))
            x_ = x - center_x
            y_ = y - center_y
            H_x = np.empty((n_order+1, n))
            H_y = np.empty((n_order+1, n))
            for n in range(0,n_order+1):
                prefactor = 1./np.sqrt(2**n*np.sqrt(np.pi)*math.factorial(n))
                n_array = np.zeros(n+1)
                n_array[n] = 1
                H_x[n] = hermite.hermval(x_/beta, n_array) * prefactor * np.exp(-(x_/beta)**2/2.)
                H_y[n] = hermite.hermval(y_/beta, n_array) * prefactor * np.exp(-(y_/beta)**2/2.)
            return H_x, H_y

        def _get_num_n(self, n_coeffs):
            """

            :param n_coeffs: number of coeffs
            :return: number of n_l of order of the shapelets
            """
            num_n = round((math.sqrt(8*n_coeffs + 1) -1)/2. +0.499)
            return int(num_n)