Source code for lenstronomy.Cosmo.cosmo_solver

__author__ = "sibirrer"

import scipy.optimize
import scipy.interpolate as interpolate
import numpy as np

from astropy.cosmology import FlatLambdaCDM
from lenstronomy.Cosmo.lens_cosmo import LensCosmo

from lenstronomy.Util.package_util import exporter

export, __all__ = exporter()


[docs] @export def cosmo2angular_diameter_distances(H_0, omega_m, z_lens, z_source): """ :param H_0: Hubble constant [km/s/Mpc] :param omega_m: dimensionless matter density at z=0 :param z_lens: deflector redshift :param z_source: source redshift :return: angular diameter distances Dd and Ds/Dds """ cosmo = FlatLambdaCDM(H0=H_0, Om0=omega_m, Ob0=0.0) lensCosmo = LensCosmo(z_lens=z_lens, z_source=z_source, cosmo=cosmo) Dd = lensCosmo.dd Ds = lensCosmo.ds Dds = lensCosmo.dds return Dd, Ds / Dds
[docs] @export def ddt2h0(ddt, z_lens, z_source, cosmo): """Converts time-delay distance to H0 for a given expansion history. :param ddt: time-delay distance in Mpc :param z_lens: deflector redshift :param z_source: source redshift :param cosmo: astropy.cosmology class instance :return: h0 value which matches the cosmology class effectively replacing the h0 value used in the creation of this class """ h0_fiducial = cosmo.H0.value lens_cosmo = LensCosmo(z_lens=z_lens, z_source=z_source, cosmo=cosmo) ddt_fiducial = lens_cosmo.ddt h0 = h0_fiducial * ddt_fiducial / ddt return h0
[docs] @export class SolverFlatLCDM(object): """Class to solve multidimensional non-linear equations to determine the cosmological parameters H0 and omega_m given the angular diameter distance relations."""
[docs] def __init__(self, z_d, z_s): self.z_d = z_d self.z_s = z_s
[docs] def F(self, x, Dd, Ds_Dds): """ :param x: array of parameters (H_0, omega_m) :return: """ [H_0, omega_m] = x omega_m = abs(omega_m) % 1 Dd_new, Ds_Dds_new = cosmo2angular_diameter_distances( H_0, omega_m, self.z_d, self.z_s ) y = np.zeros(2) y[0] = Dd - Dd_new y[1] = Ds_Dds - Ds_Dds_new return y
[docs] def solve(self, init, dd, ds_dds): x = scipy.optimize.fsolve( self.F, init, args=(dd, ds_dds), xtol=1.49012e-08, factor=0.1 ) x[1] = abs(x[1]) % 1 y = self.F(x, dd, ds_dds) if abs(y[0]) >= 0.1 or abs(y[1]) > 0.1: x = np.array([-1, -1]) return x
[docs] @export class InvertCosmo(object): """Class to do an interpolation and call the inverse of this interpolation to get H_0 and omega_m."""
[docs] def __init__(self, z_d, z_s, H0_range=None, omega_m_range=None): self.z_d = z_d self.z_s = z_s if H0_range is None: H0_range = np.linspace(10, 100, 90) if omega_m_range is None: omega_m_range = np.linspace(0.05, 1, 95) self._H0_range = H0_range self._omega_m_range = omega_m_range
def _make_interpolation(self): """Creates an interpolation grid in H_0, omega_m and computes quantities in Dd and Ds_Dds. :return: """ grid2d = np.dstack(np.meshgrid(self._H0_range, self._omega_m_range)).reshape( -1, 2 ) H0_grid = grid2d[:, 0] omega_m_grid = grid2d[:, 1] Dd_grid = np.zeros_like(H0_grid) Ds_Dds_grid = np.zeros_like(H0_grid) for i in range(len(H0_grid)): Dd, Ds_Dds = cosmo2angular_diameter_distances( H0_grid[i], omega_m_grid[i], self.z_d, self.z_s ) Dd_grid[i] = Dd Ds_Dds_grid[i] = Ds_Dds self._f_H0 = interpolate.interp2d( Dd_grid, Ds_Dds_grid, H0_grid, kind="linear", copy=False, bounds_error=False, fill_value=-1, ) print("H0 interpolation done") self._f_omega_m = interpolate.interp2d( Dd_grid, Ds_Dds_grid, omega_m_grid, kind="linear", copy=False, bounds_error=False, fill_value=0, ) print("omega_m interpolation done")
[docs] def get_cosmo(self, Dd, Ds_Dds): """Return the values of H0 and omega_m computed with an interpolation. :param Dd: flat :param Ds_Dds: float :return: """ if not hasattr(self, "_f_H0") or not hasattr(self, "_f_omega_m"): self._make_interpolation() H0 = self._f_H0(Dd, Ds_Dds) print(H0, "H0") omega_m = self._f_omega_m(Dd, Ds_Dds) Dd_new, Ds_Dds_new = cosmo2angular_diameter_distances( H0[0], omega_m[0], self.z_d, self.z_s ) if abs(Dd - Dd_new) / Dd > 0.01 or abs(Ds_Dds - Ds_Dds_new) / Ds_Dds > 0.01: return -1, -1 else: return H0[0], omega_m[0]