new charged lens class implemented

utils/lenses.py

+import numpy as np
+
+class ChargedLens(object):
+    """
+    Class to create a lens with charge
+    Note that we consider y2=0
+    
+    Initialization params:
+    -------------------------
+    qeff     : float
+                effective charge parameter
+    y1, y2   : float, float
+                dimensionless impact parameter
+    gtd_limit: float
+                max time delay gradient at images [default:4.5]
+    """
+    def __init__(self, qeff, y1, y2=0., gtd_limit=4.5):
+        """
+        params:
+        ---------------
+        qeff      : float
+                    effective charge
+        y1        : float
+                    impact parameter
+        y2        : float [default:0.0]
+                    impact parameter
+        gtd_limit : float
+                    max time delay gradient at images [default:4.5]
+        """
+        self.qeff = qeff
+        self.y1 = y1
+        self.y2 = y2
+        self.gtd_limit = gtd_limit
+        self.img, _ = self.imageLocations()
+        self.img_midx = []
+        self.img_mag = []
+        self.img_td = []
+        if len(self.img)>0:
+            self.img_td = self.timeDelay(self.img, np.array([1e-12]*len(self.img)))
+            for img in self.img:
+                self.img_midx.append(self.morseIdx(img, 1e-12))
+                self.img_mag.append(self.magnification(img,1e-12))
+        
+        
+    def potential(self, x1, x2):
+        """ Lensing potential """
+        x = np.sqrt(x1**2 + x2**2)
+        return np.log(x) + self.qeff/x
+
+    def timeDelay(self, x1, x2, minimize=False):
+        """ Lensing time delay, a.k.a. Fermat potential """
+        t_geom = ((x2 - self.y2)**2 + (x1 - self.y1)**2) / 2
+        td = t_geom - self.potential(x1, x2)
+        if minimize:
+            minima_td = np.min(self.img_td)
+            td -= minima_td
+        return td
+        
+    def alpha(self, x1, x2=0):
+        """ Deflection angle """
+        x = np.sqrt(x1**2 + x2**2)
+        a1 = x1 * (1 - self.qeff/x)/x**2
+        a2 = x2 * (1 - self.qeff/x)/x**2
+        return a1, a2
+        
+    def gradTimeDelay(self, x1, x2=0):
+        """ Gradient of time delay function """
+        a1, a2 = self.alpha(x1, x2)
+        dtdx1 = x1 - self.y1 - a1
+        dtdx2 = x2 - self.y2 - a2
+        return dtdx1, dtdx2
+
+    def critLines(self):
+        """ Tangential critical points are mapped to y=0.
+        Radial critical points are obtained here"""
+        common_term = np.complex128(9 * self.qeff + 
+                                    np.sqrt(3 * (1 + 27 * self.qeff**2)))
+        common_term = np.power(common_term, (1./3.))
+        expr1 = -1./(3**(1/3) * common_term) + (common_term/3**(2/3))
+        expr2 = (1 + 1j*np.sqrt(3))/(2*3**(1/3)*common_term) - (1 - 1j*np.sqrt(3))*common_term/(2*3**(2/3))
+        expr3 = (1 - 1j*np.sqrt(3))/(2*3**(1/3)*common_term) - (1 + 1j*np.sqrt(3))*common_term/(2*3**(2/3))
+        crit_pts = []
+        for crit_pt in [expr1, expr2, expr3]:
+            if (abs(np.imag(crit_pt)) < 1e-4):
+                crit_pts.append(np.real(crit_pt))
+        crit_pts = np.asarray(crit_pts)
+        return crit_pts
+
+    def caustic(self):
+        """ Radial caustic points """
+        crit_pts = self.critLines()
+        caustic_pts = []
+        for crit_pt in crit_pts:
+            caustic_pts.append(self.lensEquation(crit_pt))
+        caustic_pts = np.asarray(caustic_pts)
+        return caustic_pts
+        
+    def detTrace(self, x1, x2):
+        """ Returns the determinant and trace of the time delay Hessian matrix """
+        x = np.sqrt(x1**2 + x2**2)
+        term1 = 1 - 1/x**2 + self.qeff/x**3
+        term2 = 1 + 1/x**2 - 2*self.qeff/x**3
+        det = term1 * term2
+        tr = 2 - self.qeff / x**3
+        return det, tr
+
+    def magnification(self, x1, x2):
+        """ Inverse determinant of the Hessian matrix """
+        det, tr = self.detTrace(x1, x2)
+        return 1. / det
+
+    def morseIdx(self, x1im, x2im):
+        """
+        Returns the Morse Index depending on the type
+        of image, given the image location x1im, x2im.
+        """
+        det, tr = self.detTrace(x1im, x2im)
+        if det<0:
+            return 0.5 # saddle
+        elif det>0 and tr>0:
+            return 0 # minima
+        elif det>0 and tr<0:
+            return 1 # maxima
+        else:
+            return -1 # undefined
+
+    def lensEquation(self, x):
+        """ right hand side of the lens equation """
+        return x - (1/x) + (self.qeff/x**2)
+    
+    def imageLocations(self):
+        """
+        Returns the location of images in the format
+        [[im1_x1, im2_x1, ...], [im1_x2, im2_x2, ...]]
+        """
+        # for now we are doing it for y2 = 0
+        def eq1(x):
+            return x*x*self.y1
+        def eq2(x):
+            return x**3 - x + self.qeff
+        x = np.arange(-4, 4, 0.0005)
+        f1 = eq1(x)
+        f2 = eq2(x)
+        possible_img_locs = x[np.argwhere(np.diff(np.sign(f1 - f2))).flatten()]
+        gtd1, gtd2 = self.gradTimeDelay(possible_img_locs, np.zeros_like(possible_img_locs))
+        gtdnorm = np.sqrt(gtd1**2 + gtd2**2)
+        idx = np.where(gtdnorm<self.gtd_limit)[0]
+        img_locs = np.array(possible_img_locs[idx])
+        return [img_locs, np.zeros_like(img_locs)]
+
+    def imageLocationsAnalytic(self, grad_td_check=True):
+        """ Analytically obtained locations of images """
+        val = 81 * self.qeff**2 - 3 * (4 + self.y1**2) - 6 * self.qeff * self.y1 * (9 + 2 * self.y1**2)
+        if np.sign(val)==-1:
+            common_term = np.complex128(27 * self.qeff - 9 * self.y1 - 2 * self.y1**3 + 3j * np.sqrt(np.abs(val)))
+        else:
+            common_term = np.complex128(27 * self.qeff - 9 * self.y1 - 2 * self.y1**3 + 3 * np.sqrt(val))
+        common_term = np.power(common_term, (1./3.))
+        expr1 = (1/6) * (2 * self.y1 - ((2**(4/3) * (3 + self.y1**2)) / common_term) - 2**(2/3) * common_term)
+        expr2 = (1/12) * (4 * self.y1 + ((2**(4/3) * (1 + 1j * np.sqrt(3)) * (3 + self.y1**2)) / 
+                          common_term) + 2**(2/3) * (1 - 1j * np.sqrt(3)) * common_term)
+        expr3 = (1/12) * (4 * self.y1 + ((2**(4/3) * (1 - 1j * np.sqrt(3)) * (3 + self.y1**2)) / 
+                          common_term) + 2**(2/3) * (1 + 1j * np.sqrt(3)) * common_term)
+        img_locs = []
+        for imgloc in [expr1, expr2, expr3]:
+            if (abs(np.imag(imgloc)) < 1e-4):
+                if grad_td_check:
+                    gtd1, gtd2 = self.gradTimeDelay(imgloc, 0)
+                    gtdnorm = np.sqrt(gtd1**2 + gtd2**2)
+                    if gtdnorm < self.gtd_limit:
+                        img_locs.append(np.real(imgloc))
+                else:
+                    img_locs.append(np.real(imgloc))
+        img_locs = np.asarray(img_locs)
+        return img_locs
\ No newline at end of file