How to use the lenstronomy.LensModel.Solver.lens_equation_solver.LensEquationSolver function in lenstronomy

def test_solver_profile_shear(self):
        lens_model_list = ['SPEP', 'SHEAR_GAMMA_PSI']
        lensModel = LensModel(lens_model_list)

        lensEquationSolver = LensEquationSolver(lensModel)
        sourcePos_x = 0.
        sourcePos_y = 0.1
        deltapix = 0.05
        numPix = 150
        gamma = 1.98
        e1, e2 = -0.04, -0.01
        gamma_ext = np.sqrt(e1**2 + e2**2)
        kwargs_shear = {'gamma_ext': gamma_ext, 'psi_ext': 0.3}  # shear values to the source plane
        kwargs_spemd = {'theta_E': 1.66, 'gamma': gamma, 'center_x': 0.0, 'center_y': 0.0, 'e1': 0.1,
                        'e2': 0.05}  # parameters of the deflector lens model

        kwargs_lens = [kwargs_spemd, kwargs_shear]
        x_pos, y_pos = lensEquationSolver.findBrightImage(sourcePos_x, sourcePos_y, kwargs_lens, numImages=4,
                                                          min_distance=deltapix, search_window=numPix * deltapix)
        print(x_pos, y_pos, 'test positions')
        kwargs_lens_init = [{'theta_E': 1.3, 'gamma': gamma, 'e1': 0, 'e2': 0, 'center_x': 0., 'center_y': 0},

def test_logL(self):

        z_source = 1.5
        z_lens = 0.5
        cosmo = FlatLambdaCDM(H0=70, Om0=0.3, Ob0=0.)
        lensCosmo = LensCosmo(cosmo=cosmo, z_lens=z_lens, z_source=z_source)

        # make class instances for a chosen lens model type

        # chose a lens model
        lens_model_list = ['SPEP', 'SHEAR']
        # make instance of LensModel class
        lensModel = LensModel(lens_model_list=lens_model_list, cosmo=cosmo, z_lens=z_lens, z_source=z_source)
        # we require routines accessible in the LensModelExtensions class
        # make instance of LensEquationSolver to solve the lens equation
        lensEquationSolver = LensEquationSolver(lensModel=lensModel)

        # make choice of lens model

        # we chose a source position (in units angle)
        x_source, y_source = 0.02, 0.01
        # we chose a lens model
        kwargs_lens = [{'theta_E': 1., 'e1': 0.1, 'e2': 0.2, 'gamma': 2., 'center_x': 0, 'center_y': 0},
                       {'gamma1': 0.05, 'gamma2': -0.01}]

        # compute image positions and their (finite) magnifications

        # we solve for the image position(s) of the provided source position and lens model
        x_img, y_img = lensEquationSolver.image_position_from_source(kwargs_lens=kwargs_lens, sourcePos_x=x_source,
                                                                     sourcePos_y=y_source)

        point_source_list = ['LENSED_POSITION']

def test_solver_simplified(self):
        lens_model_list = ['SPEP', 'SHEAR_GAMMA_PSI']
        lensModel = LensModel(lens_model_list)

        lensEquationSolver = LensEquationSolver(lensModel)
        sourcePos_x = 0.1
        sourcePos_y = -0.1
        deltapix = 0.05
        numPix = 150
        gamma = 1.9
        gamma_ext = 0.05
        psi_ext = 0.4
        #e1, e2 = param_util.phi_gamma_ellipticity(phi=psi_ext, gamma=gamma_ext)
        kwargs_lens = [{'theta_E': 1., 'gamma': gamma, 'e1': 0.1, 'e2': -0.1, 'center_x': 0.1, 'center_y': -0.1},
                       {'gamma_ext': gamma_ext, 'psi_ext': psi_ext}]
        x_pos, y_pos = lensEquationSolver.findBrightImage(sourcePos_x, sourcePos_y, kwargs_lens, numImages=4,
                                                          min_distance=deltapix, search_window=numPix * deltapix)
        e1_new, e2_new = param_util.shear_polar2cartesian(phi=0., gamma=gamma_ext + 0.1)
        kwargs_lens_init = [{'theta_E': 1.3, 'gamma': gamma, 'e1': 0., 'e2': 0., 'center_x': 0., 'center_y': 0},
                            {'gamma_ext': gamma_ext + 0.1, 'psi_ext': 0}]
        solver = Solver4Point(lensModel, solver_type='PROFILE_SHEAR')

R_slit = 3.8
        dR_slit = 1.
        aperture_type = 'slit'
        kwargs_aperture = {'aperture_type': aperture_type, 'center_ra': 0, 'width': dR_slit, 'length': R_slit,
                           'angle': 0, 'center_dec': 0}
        psf_fwhm = 0.7
        kwargs_seeing = {'psf_type': 'GAUSSIAN', 'fwhm': psf_fwhm}

        TDCosmography(z_lens, z_source, kwargs_model)
        from astropy.cosmology import FlatLambdaCDM
        cosmo = FlatLambdaCDM(H0=70, Om0=0.3, Ob0=0.05)
        self.td_cosmo = TDCosmography(z_lens, z_source, kwargs_model, cosmo_fiducial=cosmo, lens_model_kinematics_bool=None,
                                      kwargs_aperture = kwargs_aperture, kwargs_seeing = kwargs_seeing,
                                      light_model_kinematics_bool=None)
        self.lens = LensModel(lens_model_list=['SIE'], cosmo=cosmo, z_lens=z_lens, z_source=z_source)
        self.solver = LensEquationSolver(lensModel=self.lens)

        self.kwargs_lens = [{'theta_E': 1, 'e1': 0.1, 'e2': -0.2, 'center_x': 0, 'center_y': 0}]
        source_x, source_y = 0, 0.05
        image_x, image_y = self.solver.image_position_from_source(source_x, source_y, self.kwargs_lens, min_distance=0.1, search_window=10)
        self.kwargs_ps = [{'ra_image': image_x, 'dec_image': image_y}]
        self.image_x, self.image_y = image_x, image_y

self.kwargs_lens_light = [kwargs_sersic]
        lens_light_model_class = LightModel(light_model_list=lens_light_model_list)
        source_model_list = ['SERSIC_ELLIPSE']
        self.kwargs_source = [kwargs_sersic_ellipse]
        source_model_class = LightModel(light_model_list=source_model_list)
        self.kwargs_ps = [{'ra_source': 0.0001, 'dec_source': 0.0,
                           'source_amp': 1.}]  # quasar point source position in the source plane and intrinsic brightness
        point_source_class = PointSource(point_source_type_list=['SOURCE_POSITION'], fixed_magnification_list=[True])
        kwargs_numerics = {'subgrid_res': 2, 'psf_subgrid': True}
        imageModel = ImageModel(data_class, psf_class, lens_model_class, source_model_class, lens_light_model_class,
                                point_source_class, kwargs_numerics=kwargs_numerics)
        image_sim = sim_util.simulate_simple(imageModel, self.kwargs_lens, self.kwargs_source,
                                         self.kwargs_lens_light, self.kwargs_ps)
        data_class.update_data(image_sim)
        kwargs_data['image_data'] = image_sim
        self.solver = LensEquationSolver(lensModel=lens_model_class)
        idex_lens_1 = [0, 1]
        idex_lens_2 = [2, 3]
        multi_band_list = [[kwargs_data, kwargs_psf, kwargs_numerics, {'index_lens_model_list': idex_lens_1}], [kwargs_data, kwargs_psf, kwargs_numerics, {'index_lens_model_list': idex_lens_2}]]
        lens_model_list_joint = lens_model_list + lens_model_list
        self.kwargs_lens_joint = self.kwargs_lens + self.kwargs_lens
        self.imageModel = MultiFrame(multi_band_list, lens_model_list_joint, source_model_class, lens_light_model_class, point_source_class)

def test_solver_multiplane(self):
        lens_model_list = ['SPEP', 'SHEAR', 'SIS']
        lensModel = LensModel(lens_model_list, z_source=1, lens_redshift_list=[0.5, 0.5, 0.3], multi_plane=True)

        lensEquationSolver = LensEquationSolver(lensModel)
        sourcePos_x = 0.01
        sourcePos_y = -0.01
        deltapix = 0.05
        numPix = 150
        gamma = 1.96
        gamma1, gamma2 = 0.01, 0.01
        kwargs_shear = {'gamma1': gamma1, 'gamma2': gamma2}  # gamma_ext: shear strength, psi_ext: shear angel (in radian)
        kwargs_spemd = {'theta_E': 1., 'gamma': gamma, 'center_x': 0, 'center_y': 0, 'e1': 0.2, 'e2': 0.03}
        kwargs_sis = {'theta_E': .1, 'center_x': 1, 'center_y': 0}
        kwargs_lens = [kwargs_spemd, kwargs_shear, kwargs_sis]
        x_pos, y_pos = lensEquationSolver.findBrightImage(sourcePos_x, sourcePos_y, kwargs_lens, numImages=4,
                                                          min_distance=deltapix, search_window=numPix * deltapix)
        print(x_pos, y_pos, 'test positions')
        kwargs_lens_init = [{'theta_E': 1.3, 'gamma': gamma, 'e1': 0.1, 'e2': 0, 'center_x': 0., 'center_y': 0},
                            {'gamma1': gamma1, 'gamma2': gamma2}, {'theta_E': .1, 'center_x': 1, 'center_y': 0}]
        solver = Solver4Point(lensModel, solver_type='PROFILE')

def test_example(self):
        lens_model_list = ['SPEP', 'SHEAR']
        lensModel = LensModel(lens_model_list)

        lensEquationSolver = LensEquationSolver(lensModel)
        sourcePos_x = 0.03
        sourcePos_y = 0.0
        min_distance = 0.05
        search_window = 10
        gamma = 2.
        gamma1, gamma2 = -0.04, -0.1
        kwargs_shear = {'gamma1': gamma1, 'gamma2': gamma2}  # shear values to the source plane
        kwargs_spemd = {'theta_E': 1., 'gamma': gamma, 'center_x': 0.0, 'center_y': 0.0, 'e1': 0.01,
                        'e2': 0.05}  # parameters of the deflector lens model

        kwargs_lens = [kwargs_spemd, kwargs_shear]
        x_pos, y_pos = lensEquationSolver.image_position_from_source(sourcePos_x, sourcePos_y, kwargs_lens,
                                                                     min_distance=min_distance,
                                                                     search_window=search_window,
                                                                     precision_limit=10 ** (-10), num_iter_max=10,
                                                                     arrival_time_sort=True)

self.kwargs_lens_light = [kwargs_sersic, kwargs_sersic]
        lens_light_model_class = LightModel(light_model_list=['SERSIC'])
        source_model_list = ['SERSIC_ELLIPSE', 'SERSIC_ELLIPSE']
        self.kwargs_source = [kwargs_sersic_ellipse, kwargs_sersic_ellipse]
        source_model_class = LightModel(light_model_list=['SERSIC_ELLIPSE'])
        self.kwargs_ps = [{'ra_source': 0.0001, 'dec_source': 0.0,
                           'source_amp': 1.}]  # quasar point source position in the source plane and intrinsic brightness
        point_source_class = PointSource(point_source_type_list=['SOURCE_POSITION'], fixed_magnification_list=[True])
        kwargs_numerics = {'subgrid_res': 2, 'psf_subgrid': True}
        imageModel = ImageModel(data_class, psf_class, lens_model_class, source_model_class, lens_light_model_class,
                                point_source_class, kwargs_numerics=kwargs_numerics)
        image_sim = sim_util.simulate_simple(imageModel, self.kwargs_lens, [kwargs_sersic_ellipse],
                                             [kwargs_sersic], self.kwargs_ps)
        data_class.update_data(image_sim)
        kwargs_data['image_data'] = image_sim
        self.solver = LensEquationSolver(lensModel=lens_model_class)
        kwargs_model_bool = {'index_source_light_model': [0], 'index_lens_light_model': [0]}
        multi_band_list = [[kwargs_data, kwargs_psf, kwargs_numerics, kwargs_model_bool]]
        self.imageModel = MultiBandMultiModel(multi_band_list, lens_model_class, source_model_list, lens_light_model_list, point_source_class)

def test_decoupling(self):
        lens_model_list = ['SPEP', 'SIS']
        lensModel = LensModel(lens_model_list)
        solver = Solver4Point(lensModel)
        solver_decoupled = Solver4Point(lensModel)
        lensEquationSolver = LensEquationSolver(lensModel)
        sourcePos_x = 0.1
        sourcePos_y = -0.1
        deltapix = 0.05
        numPix = 150
        gamma = 1.9
        phi_G, q = 0.5, 0.8
        e1, e2 = param_util.phi_q2_ellipticity(phi_G, q)
        kwargs_lens = [{'theta_E': 1., 'gamma': gamma, 'e1': e1, 'e2': e2, 'center_x': 0.1, 'center_y': -0.1},
                       {'theta_E': 0.1, 'center_x': 0.5, 'center_y': 0}]
        x_pos, y_pos = lensEquationSolver.findBrightImage(sourcePos_x, sourcePos_y, kwargs_lens, numImages=4, min_distance=deltapix, search_window=numPix*deltapix)
        phi_G, q = 1.5, 0.9
        e1, e2 = param_util.phi_q2_ellipticity(phi_G, q)
        kwargs_lens_init = [{'theta_E': 1.3, 'gamma': gamma, 'e1': e1, 'e2': e2, 'center_x': 0., 'center_y': 0}, {'theta_E': 0.1, 'center_x': 0.5, 'center_y': 0}]
        kwargs_lens_new, accuracy = solver.constraint_lensmodel(x_pos, y_pos, kwargs_lens_init)
        kwargs_lens_new_2, accuracy = solver_decoupled.constraint_lensmodel(x_pos, y_pos, kwargs_lens_init)
        print(kwargs_lens_new_2)

def _find_point_sources(self, kwargs_options, kwargs_lens, kwargs_else):
        lensModel = LensModel(kwargs_options.get('lens_model_list', []))
        imPos = LensEquationSolver(lensModel)
        if kwargs_options.get('point_source', False):
            min_distance = 0.05
            search_window = 10
            sourcePos_x = kwargs_else['sourcePos_x']
            sourcePos_y = kwargs_else['sourcePos_y']
            x_mins, y_mins = imPos.image_position_from_source(sourcePos_x, sourcePos_y, kwargs_lens, min_distance=min_distance, search_window=search_window)
            n = len(x_mins)
            mag_list = np.zeros(n)
            for i in range(n):
                mag = lensModel.magnification(x_mins[i], y_mins[i], kwargs_lens)
                mag_list[i] = abs(mag)
            kwargs_else['ra_pos'] = x_mins
            kwargs_else['dec_pos'] = y_mins
            kwargs_else['point_amp'] = mag_list * kwargs_else['quasar_amp']
        return kwargs_else