added image finder code

utils/image_finder.py

+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Created on Mon Apr  4 10:43:42 2022
+
+@author: uddeepta
+"""
+import numpy as np
+from scipy.optimize import minimize
+
+
+def find_images(lens, method='L', **kwargs):
+    if method == 'T':
+        x1_, x2_ = imageLocationsTriangleMapping(lens, **kwargs)
+    elif method == 'S':
+        x1_, x2_ = imageLocationsStochastic(lens, **kwargs)
+    elif method == 'L':
+        x1_, x2_ = imageLocationsLenstronomy(lens, **kwargs)
+    return x1_, x2_
+
+
+def rayShooting(lens, x1, x2):
+    """
+    Returns the maps of image at coordinates (x1, x2) to source position (inverse deflection)
+    """
+    dx1, dx2 = lens.alpha(x1, x2)
+    return x1 - dx1, x2 - dx2
+
+
+def imageLocationsTriangleMapping(lens, search_window=5, npix=2048):
+    """
+    Method of finding images using *Triangle mapping* 
+    (see Bartelmann 2003, Schneider et al. 1992)
+    
+    For each pixel on the lens plane, build two triangles.
+    Using ray-tracing, these are mapped onto the source plane into 
+    the other two triangles. Compute the distances of the vertices 
+    of the triangles on the source plane from the source and check 
+    using cross-products if the source is inside one of the two triangles
+    
+    params:
+    ---------------
+    lens: an object of class lens
+    search_window : maxima of the search domain
+    npix : number of pixels
+    """
+    ys1 = lens.y1
+    ys2 = lens.y2
+
+    # create a mesh in the search region: 
+    x = np.linspace(-abs(search_window), abs(search_window), npix)
+    grid_pixel = x[1] - x[0]
+    x1, x2 = np.meshgrid(x, x)
+
+    x = np.sqrt(x1 * x1 + x2 * x2)
+
+    # ray-trace the deflector grid onto the source plane
+    y1, y2 = rayShooting(lens, x1, x2)
+
+    # convert to pixel units
+    xray = y1 / grid_pixel + (len(x) - 1) / 2.0
+    yray = y2 / grid_pixel + (len(x) - 1) / 2.0
+
+    y1s = ys1 / grid_pixel + (len(x) - 1) / 2.0
+    y2s = ys2 / grid_pixel + (len(x) - 1) / 2.0
+
+    # shift the maps by one pixel
+    xray1 = np.roll(xray, 1, axis=1)
+    xray2 = np.roll(xray1, 1, axis=0)
+    xray3 = np.roll(xray2, -1, axis=1)
+    yray1 = np.roll(yray, 1, axis=1)
+    yray2 = np.roll(yray1, 1, axis=0)
+    yray3 = np.roll(yray2, -1, axis=1)
+
+    x1 = y1s - xray
+    y1 = y2s - yray
+    x2 = y1s - xray1
+    y2 = y2s - yray1
+    x3 = y1s - xray2
+    y3 = y2s - yray2
+    x4 = y1s - xray3
+    y4 = y2s - yray3
+
+    prod12 = x1 * y2 - x2 * y1
+    prod23 = x2 * y3 - x3 * y2
+    prod31 = x3 * y1 - x1 * y3
+    prod13 = -prod31
+    prod34 = x3 * y4 - x4 * y3
+    prod41 = x4 * y1 - x1 * y4
+
+    image = np.zeros(xray.shape)
+    image[((np.sign(prod12) == np.sign(prod23)) &
+           (np.sign(prod23) == np.sign(prod31)))] = 1
+    image[((np.sign(prod13) == np.sign(prod34)) &
+           (np.sign(prod34) == np.sign(prod41)))] = 2
+
+    # first kind of images (first triangle)
+    images1 = np.argwhere(image == 1)
+    xi_images_ = images1[:, 1]
+    yi_images_ = images1[:, 0]
+    xi_images = xi_images_[(xi_images_ > 0) & (yi_images_ > 0)]
+    yi_images = yi_images_[(xi_images_ > 0) & (yi_images_ > 0)]
+
+    # compute the weights
+    w = np.array([1. / np.sqrt(x1[xi_images, yi_images] ** 2 +
+                               y1[xi_images, yi_images] ** 2),
+                  1. / np.sqrt(x2[xi_images, yi_images] ** 2 +
+                               y2[xi_images, yi_images] ** 2),
+                  1. / np.sqrt(x3[xi_images, yi_images] ** 2 +
+                               y3[xi_images, yi_images] ** 2)])
+    xif1, yif1 = _refineImagePositions(xi_images, yi_images, w, 1)
+
+    # second kind of images (second triangle)
+    images1 = np.argwhere(image == 2)
+    xi_images_ = images1[:, 1]
+    yi_images_ = images1[:, 0]
+    xi_images = xi_images_[(xi_images_ > 0) & (yi_images_ > 0)]
+    yi_images = yi_images_[(xi_images_ > 0) & (yi_images_ > 0)]
+
+    # compute the weights
+    w = np.array([1. / np.sqrt(x1[xi_images, yi_images] ** 2 +
+                               y1[xi_images, yi_images] ** 2),
+                  1. / np.sqrt(x3[xi_images, yi_images] ** 2 +
+                               y3[xi_images, yi_images] ** 2),
+                  1. / np.sqrt(x4[xi_images, yi_images] ** 2 +
+                               y4[xi_images, yi_images] ** 2)])
+    xif2, yif2 = _refineImagePositions(xi_images, yi_images, w, 2)
+
+    xi = np.concatenate([xif1, xif2])
+    yi = np.concatenate([yif1, yif2])
+    xi = (xi - 1 - (len(x) - 1) / 2.0) * grid_pixel
+    yi = (yi - 1 - (len(x) - 1) / 2.0) * grid_pixel
+    return xi, yi
+
+
+def _refineImagePositions(x, y, w, typ):
+    """
+    Image positions are calculated as the weighted mean of the positions 
+    of the triangle vertices. The weights are the distances between the vertices 
+    mapped onto the source plane and the source position.
+    """
+    if typ == 2:
+        xp = np.array([x, x + 1, x + 1])
+        yp = np.array([y, y, y + 1])
+    else:
+        xp = np.array([x, x + 1, x])
+        yp = np.array([y, y + 1, y + 1])
+    xi = np.zeros(x.size)
+    yi = np.zeros(y.size)
+    for i in range(x.size):
+        xi[i] = (xp[:, i] / w[:, i]).sum() / (1. / w[:, i]).sum()
+        yi[i] = (yp[:, i] / w[:, i]).sum() / (1. / w[:, i]).sum()
+    return xi, yi
+
+
+def imageLocationsStochastic(lens, search_window=10, precision_limit=1e-10, x_center=0, y_center=0, num_random=1000):
+    """
+    Solves the lens equation stochastic with the scipy minimization routine 
+    on the quadratic distance between the backwards ray-shooted proposed image position and the source position.
+    Credits to Giulia Pagano
+
+    """
+    source_x, source_y = lens.y1, lens.y2
+    x_solve, y_solve = [], []
+    for i in range(num_random):
+        x_init = np.random.uniform(-search_window / 2., search_window / 2) + x_center
+        y_init = np.random.uniform(-search_window / 2., search_window / 2) + y_center
+        xinitial = np.array([x_init, y_init])
+        result = minimize(_root, xinitial, args=(source_x, source_y, lens), tol=precision_limit ** 2,
+                          method='Nelder-Mead')
+        if _root(result.x, source_x, source_y, lens) < precision_limit ** 2:
+            x_solve.append(result.x[0])
+            y_solve.append(result.x[1])
+    x_mins, y_mins = _findOverlap(x_solve, y_solve, precision_limit)
+    return x_mins, y_mins
+
+
+def _root(x, source_x, source_y, lens):
+    """
+    Returns square distance between ray-traced image position and given source position
+    """
+    x_, y_ = x
+    beta_x, beta_y = rayShooting(lens, x_, y_)
+    return (beta_x - source_x) ** 2 + (beta_y - source_y) ** 2
+
+
+def _findOverlap(x_mins, y_mins, min_distance):
+    """
+    finds overlapping solutions, deletes multiples and deletes non-solutions
+    and if it is not a solution, deleted as well
+    """
+    n = len(x_mins)
+    idex = []
+    for i in range(n):
+        if i == 0:
+            pass
+        else:
+            for j in range(0, i):
+                if abs(x_mins[i] - x_mins[j] < min_distance and abs(y_mins[i] - y_mins[j]) < min_distance):
+                    idex.append(i)
+                    break
+    x_mins = np.delete(x_mins, idex, axis=0)
+    y_mins = np.delete(y_mins, idex, axis=0)
+    return x_mins, y_mins
+
+
+def imageLocationsLenstronomy(lens, min_distance=0.01, search_window=10, precision_limit=1e-10, num_iter_max=100,
+                              initial_guess_cut=True, x_center=0, y_center=0, num_random=0,
+                              non_linear=False, magnification_limit=None):
+    """
+    Finds image position  given source position and lens model. The solver first does a grid search in the
+    lens plane, and the grid points that are closest to the supplied source position are fed to a
+    specialized gradient-based root finder that finds the exact solutions. Works with all lens models.
+
+    """
+    # find pixels in the image plane possibly hosting a solution of the lens equation, related source distances and
+    # pixel width
+    x_mins, y_mins, delta_map, pixel_width = _candidate_solutions(lens, min_distance, search_window, x_center, y_center)
+    if len(x_mins) < 1:
+        return x_mins, y_mins
+    if initial_guess_cut:
+        mag = np.abs(lens.magnification(x_mins, y_mins))
+        mag[mag < 1] = 1
+        x_mins = x_mins[delta_map <= min_distance * mag * 5]
+        y_mins = y_mins[delta_map <= min_distance * mag * 5]
+
+    x_mins = np.append(x_mins, np.random.uniform(low=-search_window / 2 + x_center,
+                                                 high=search_window / 2 + x_center, size=num_random))
+    y_mins = np.append(y_mins, np.random.uniform(low=-search_window / 2 + y_center,
+                                                 high=search_window / 2 + y_center, size=num_random))
+    # iterative solving of the lens equation for the selected grid points
+    # print("Candidates:", x_mins.shape, y_mins.shape)
+    x_mins, y_mins, solver_precision = _find_gradient_descent(lens, x_mins, y_mins, precision_limit, num_iter_max,
+                                                              min_distance=min_distance, non_linear=non_linear)
+    # only select iterative results that match the precision limit
+    x_mins = x_mins[solver_precision <= precision_limit]
+    y_mins = y_mins[solver_precision <= precision_limit]
+    # find redundant solutions within the min_distance criterion
+    x_mins, y_mins = _findOverlap(x_mins, y_mins, min_distance)
+
+    if magnification_limit is not None:
+        mag = np.abs(lens.magnification(x_mins, y_mins))
+        x_mins = x_mins[mag >= magnification_limit]
+        y_mins = y_mins[mag >= magnification_limit]
+    return x_mins, y_mins
+
+
+def _candidate_solutions(lens, min_distance=0.1, search_window=10, x_center=0, y_center=0):
+    """
+    finds pixels in the image plane possibly hosting a solution of the lens equation,
+    for the given source position and lens model
+    """
+    # compute number of pixels to cover the search window with the required min_distance
+    numPix = int(round(search_window / min_distance) + 0.5)
+    x_grid, y_grid = _make_grid(numPix, min_distance)
+    x_grid += x_center
+    y_grid += y_center
+    # ray-shoot to find the relative distance to the required source position for each grid point
+    x_mapped, y_mapped = rayShooting(lens, x_grid, y_grid)
+    absmapped = _displaceAbs(lens, x_mapped, y_mapped)
+    # select minima in the grid points and select grid points that do not deviate more than the
+    # width of the grid point to a solution of the lens equation
+    x_mins, y_mins, delta_map = _neighborSelect(absmapped, x_grid, y_grid)
+    # pixel width
+    pixel_width = x_grid[1] - x_grid[0]
+    return x_mins, y_mins, delta_map, pixel_width
+
+
+def _find_gradient_descent(lens, x_min, y_min, precision_limit=1e-9,
+                           num_iter_max=200, min_distance=0.01, non_linear=False):
+    """
+    given a 'good guess' of a solution of the lens equation (expected image position given a fixed source position)
+    this routine iteratively performs a ray-tracing with second order correction (effectively gradient decent) to find
+    a precise solution to the lens equation.
+    Return: x_position array, y_position array, error in the source plane array
+    """
+    num_candidates = len(x_min)
+    x_mins = np.zeros(num_candidates)
+    y_mins = np.zeros(num_candidates)
+    solver_precision = np.zeros(num_candidates)
+    for i in range(len(x_min)):
+        x_guess, y_guess, delta, l = _solve_single_proposal(lens, x_min[i], y_min[i], precision_limit, num_iter_max,
+                                                            max_step=min_distance, non_linear=non_linear)
+
+        x_mins[i] = x_guess
+        y_mins[i] = y_guess
+        solver_precision[i] = delta
+    return x_mins, y_mins, solver_precision
+
+
+def _solve_single_proposal(lens, x_guess, y_guess, precision_limit, num_iter_max,
+                           max_step, non_linear=False):
+    """
+    gradient decent solution of a single proposed starting point (close to a true solution)
+    return: x_position, y_position, error in the source plane, steps required (for gradient decent)
+    """
+    source_x, source_y = lens.y1, lens.y2
+    l = 0
+    if non_linear:
+        xinitial = np.array([x_guess, y_guess])
+        result = minimize(_root, xinitial, tol=precision_limit ** 2,
+                          method='Nelder-Mead')
+        delta = _root(result.x, source_x, source_y, lens)
+        x_guess, y_guess = result.x[0], result.x[1]
+
+    else:
+        x_mapped, y_mapped = rayShooting(lens, x_guess, y_guess)
+        delta = np.sqrt((x_mapped - source_x) ** 2 + (y_mapped - source_y) ** 2)
+
+        while delta > precision_limit and l < num_iter_max:
+            x_mapped, y_mapped = rayShooting(lens, x_guess, y_guess)
+            delta = np.sqrt((x_mapped - source_x) ** 2 + (y_mapped - source_y) ** 2)
+            f_xx, f_xy, f_yx, f_yy = lens.hessian(x_guess, y_guess)
+            DistMatrix = np.array([[1 - f_yy, f_yx], [f_xy, 1 - f_xx]])
+            det, tr = lens.detTraceInverseMagnificationMatrix(x_guess, y_guess)
+            deltaVec = np.array([x_mapped - source_x, y_mapped - source_y])
+            image_plane_vector = DistMatrix.dot(deltaVec) / det
+            dist = np.sqrt(image_plane_vector[0] ** 2 + image_plane_vector[1] ** 2)
+            if dist > max_step:
+                image_plane_vector *= max_step / dist
+            x_guess, y_guess, delta, l = _gradient_step(lens, x_guess, y_guess, delta,
+                                                        image_plane_vector, l, num_iter_max)
+    return x_guess, y_guess, delta, l
+
+
+def _gradient_step(lens, x_guess, y_guess, delta_init, image_plane_vector,
+                   iter_num, num_iter_max):
+    """
+    Returns updated image position in x, updated image position in y, 
+    updated precision in the source plane, total iterations done after this call
+    """
+    source_x, source_y = lens.y1, lens.y2
+    x_new = x_guess - image_plane_vector[0]
+    y_new = y_guess - image_plane_vector[1]
+    x_mapped, y_mapped = rayShooting(lens, x_new, y_new)
+    delta_new = np.sqrt((x_mapped - source_x) ** 2 + (y_mapped - source_y) ** 2)
+    iter_num += 1
+    if delta_new > delta_init:
+        if num_iter_max < iter_num:
+            return x_guess, y_guess, delta_init, iter_num
+        else:
+            # if the new proposal is worse than the previous one, it randomly draws a new proposal in a different
+            # direction and tries again
+            image_plane_vector[0] *= np.random.normal(loc=0, scale=0.5)
+            image_plane_vector[1] *= np.random.normal(loc=0, scale=0.5)
+            return _gradient_step(lens, x_guess, y_guess, delta_init, image_plane_vector, iter_num, num_iter_max)
+    else:
+        return x_new, y_new, delta_new, iter_num
+
+
+def _displaceAbs(lens, x, y):
+    """
+    calculates a grid of distances to the observer in angle
+    returns an array of displacement
+    """
+    x_mapped = x - lens.y1
+    y_mapped = y - lens.y2
+    absmapped = np.sqrt(x_mapped ** 2 + y_mapped ** 2)
+    return absmapped
+
+
+def _make_grid(numPix, deltapix, subgrid_res=1, left_lower=False):
+    """
+    creates pixel grid (in 1d arrays of x- and y- positions)
+    returns x, y position information in two 1d arrays
+    """
+
+    # Check numPix is an integer, or 2-sequence of integers
+    if isinstance(numPix, (tuple, list, np.ndarray)):
+        assert len(numPix) == 2
+        if any(x != round(x) for x in numPix):
+            raise ValueError("numPix contains non-integers: %s" % numPix)
+        numPix = np.asarray(numPix, dtype=int)
+    else:
+        if numPix != round(numPix):
+            raise ValueError("Attempt to specify non-int numPix: %s" % numPix)
+        numPix = np.array([numPix, numPix], dtype=int)
+
+    # Super-resolution sampling
+    numPix_eff = (numPix * subgrid_res).astype(int)
+    deltapix_eff = deltapix / float(subgrid_res)
+
+    # Compute unshifted grids.
+    # X values change quickly, Y values are repeated many times
+    x_grid = np.tile(np.arange(numPix_eff[0]), numPix_eff[1]) * deltapix_eff
+    y_grid = np.repeat(np.arange(numPix_eff[1]), numPix_eff[0]) * deltapix_eff
+
+    if left_lower is True:
+        # Shift so (0, 0) is in the "lower left"
+        # Note this does not shift when subgrid_res = 1
+        shift = -1. / 2 + 1. / (2 * subgrid_res) * np.array([1, 1])
+    else:
+        # Shift so (0, 0) is centered
+        shift = deltapix_eff * (numPix_eff - 1) / 2
+
+    return x_grid - shift[0], y_grid - shift[1]
+
+
+def _neighborSelect(a, x, y):
+    """
+    returns an array of indices of local minima, values of those minima
+    """
+    dim = int(np.sqrt(len(a)))
+    values = []
+    x_mins = []
+    y_mins = []
+    for i in range(dim + 1, len(a) - dim - 1):
+        if (a[i] < a[i - 1]
+                and a[i] < a[i + 1]
+                and a[i] < a[i - dim]
+                and a[i] < a[i + dim]
+                and a[i] < a[i - (dim - 1)]
+                and a[i] < a[i - (dim + 1)]
+                and a[i] < a[i + (dim - 1)]
+                and a[i] < a[i + (dim + 1)]):
+            if (a[i] < a[(i - 2 * dim - 1) % dim ** 2]
+                    and a[i] < a[(i - 2 * dim + 1) % dim ** 2]
+                    and a[i] < a[(i - dim - 2) % dim ** 2]
+                    and a[i] < a[(i - dim + 2) % dim ** 2]
+                    and a[i] < a[(i + dim - 2) % dim ** 2]
+                    and a[i] < a[(i + dim + 2) % dim ** 2]
+                    and a[i] < a[(i + 2 * dim - 1) % dim ** 2]
+                    and a[i] < a[(i + 2 * dim + 1) % dim ** 2]
+                    and a[i] < a[(i + 2 * dim) % dim ** 2]
+                    and a[i] < a[(i - 2 * dim) % dim ** 2]
+                    and a[i] < a[(i - 2) % dim ** 2]
+                    and a[i] < a[(i + 2) % dim ** 2]):
+                x_mins.append(x[i])
+                y_mins.append(y[i])
+                values.append(a[i])
+    return np.array(x_mins), np.array(y_mins), np.array(values)