Cosmology

class FLRW:
    """
    Friedmann-Lemaître-Robertson-Walker cosmology base class.
    
    Parameters:
    - H0: Hubble constant at z=0
    - Om0: Omega matter at z=0  
    - Ode0: Omega dark energy at z=0
    - Tcmb0: CMB temperature at z=0
    - Neff: effective number of neutrino species
    - m_nu: neutrino mass
    - Ob0: Omega baryon at z=0
    - name: cosmology name
    """
    def __init__(self, H0, Om0, Ode0=None, Tcmb0=2.725*u.K, Neff=3.04, 
                 m_nu=0.0*u.eV, Ob0=None, name=None, meta=None): ...
    
    def age(self, z):
        """
        Age of universe at redshift z.
        
        Parameters:
        - z: redshift
        
        Returns:
        Quantity: age in time units
        """
    
    def lookback_time(self, z):
        """Lookback time to redshift z."""
    
    def comoving_distance(self, z):
        """Comoving distance to redshift z."""
    
    def angular_diameter_distance(self, z):
        """Angular diameter distance to redshift z."""
    
    def luminosity_distance(self, z):
        """Luminosity distance to redshift z."""
    
    def comoving_volume(self, z):
        """Comoving volume out to redshift z."""
    
    def differential_comoving_volume(self, z):
        """Differential comoving volume at redshift z."""
    
    def angular_diameter_distance_z1z2(self, z1, z2):
        """Angular diameter distance between two redshifts."""
    
    def absorption_distance(self, z):
        """Absorption distance for IGM studies."""
    
    def distmod(self, z):
        """Distance modulus to redshift z."""
    
    def kpc_comoving_per_arcmin(self, z):
        """Comoving kpc per arcminute at redshift z."""
    
    def kpc_proper_per_arcmin(self, z):
        """Proper kpc per arcminute at redshift z."""
    
    def arcsec_per_kpc_comoving(self, z):
        """Arcseconds per comoving kpc at redshift z."""
    
    def arcsec_per_kpc_proper(self, z):
        """Arcseconds per proper kpc at redshift z."""
    
    def critical_density(self, z):
        """Critical density at redshift z."""
    
    def H(self, z):
        """Hubble parameter at redshift z."""
    
    def scale_factor(self, z):
        """Scale factor at redshift z."""
    
    def Tcmb(self, z):
        """CMB temperature at redshift z."""

class LambdaCDM(FLRW):
    """
    Lambda-CDM cosmology with curvature.
    
    Parameters:
    - H0: Hubble constant
    - Om0: matter density parameter
    - Ode0: dark energy density parameter (default: 1 - Om0 - Ok0)
    - Ok0: curvature density parameter (default: 0)
    """
    def __init__(self, H0, Om0, Ode0=None, **kwargs): ...

class FlatLambdaCDM(FLRW):
    """Flat Lambda-CDM cosmology (Ok0 = 0)."""
    def __init__(self, H0, Om0, **kwargs): ...

class wCDM(FLRW):
    """CDM cosmology with constant dark energy equation of state w."""
    def __init__(self, H0, Om0, Ode0, w0, **kwargs): ...

class FlatwCDM(FLRW):
    """Flat wCDM cosmology."""
    def __init__(self, H0, Om0, w0, **kwargs): ...

class w0waCDM(FLRW):
    """CDM with time-varying dark energy: w(z) = w0 + wa * z/(1+z)."""
    def __init__(self, H0, Om0, Ode0, w0, wa, **kwargs): ...

class Flatw0waCDM(FLRW):
    """Flat w0waCDM cosmology."""
    def __init__(self, H0, Om0, w0, wa, **kwargs): ...

# Planck mission results
Planck18: FlatLambdaCDM  # Planck 2018 results
Planck15: FlatLambdaCDM  # Planck 2015 results  
Planck13: FlatLambdaCDM  # Planck 2013 results

# WMAP mission results
WMAP9: FlatLambdaCDM     # WMAP 9-year results
WMAP7: FlatLambdaCDM     # WMAP 7-year results
WMAP5: FlatLambdaCDM     # WMAP 5-year results

# Example parameter values (Planck 2018)
# H0 = 67.36 km/s/Mpc
# Om0 = 0.3153
# Ode0 = 0.6847  
# Tcmb0 = 2.7255 K
# Ob0 = 0.04930

def z_at_value(func, fval, zmin=1e-8, zmax=1000, ztol=1e-8, maxfun=500):
    """
    Find redshift at which cosmological function equals specified value.
    
    Parameters:
    - func: cosmological function (e.g., cosmo.age)
    - fval: target function value
    - zmin, zmax: redshift search range
    - ztol: redshift tolerance
    - maxfun: maximum function evaluations
    
    Returns:
    float: redshift at which func(z) = fval
    """

def cosmology_equal(cosmo1, cosmo2, format=False):
    """
    Test equality of two cosmologies.
    
    Parameters:
    - cosmo1, cosmo2: cosmology instances
    - format: return detailed comparison
    
    Returns:
    bool or dict: equality result
    """

# Default cosmology management
from astropy.cosmology import default_cosmology

@default_cosmology.getter
def get_default_cosmology():
    """Get current default cosmology."""

@default_cosmology.setter  
def set_default_cosmology(cosmo):
    """Set default cosmology."""

# Available cosmologies
from astropy.cosmology import available
available: list  # List of available standard cosmologies

class Parameter:
    """
    Cosmological parameter with units and metadata.
    
    Parameters:
    - value: parameter value
    - unit: astropy unit
    - equivalencies: unit equivalencies
    - fvalidate: validation function
    - doc: parameter documentation
    """
    def __init__(self, doc='', unit=None, equivalencies=[], fvalidate='default'): ...
    
    @property
    def value(self):
        """Parameter value."""
    
    @property
    def unit(self):
        """Parameter unit."""

from astropy.cosmology import Planck18
import astropy.units as u
import numpy as np

# Use standard Planck 2018 cosmology
cosmo = Planck18

# Basic distance calculations
z = 1.0  # redshift
d_L = cosmo.luminosity_distance(z)
d_A = cosmo.angular_diameter_distance(z)
d_C = cosmo.comoving_distance(z)

print(f"At z = {z}:")
print(f"  Luminosity distance: {d_L}")
print(f"  Angular diameter distance: {d_A}")
print(f"  Comoving distance: {d_C}")

# Time calculations
age_z1 = cosmo.age(z)
age_now = cosmo.age(0)
lookback = cosmo.lookback_time(z)

print(f"Age at z=1: {age_z1.to(u.Gyr)}")
print(f"Age now: {age_now.to(u.Gyr)}")
print(f"Lookback time: {lookback.to(u.Gyr)}")

from astropy.cosmology import FlatLambdaCDM, wCDM
import astropy.units as u

# Create custom flat Lambda-CDM model
my_cosmo = FlatLambdaCDM(H0=70*u.km/u.s/u.Mpc, Om0=0.3, 
                         Tcmb0=2.725*u.K, name="My Cosmology")

# Create model with dark energy equation of state
phantom_cosmo = wCDM(H0=70*u.km/u.s/u.Mpc, Om0=0.3, 
                     Ode0=0.7, w0=-1.2, name="Phantom DE")

# Compare distances
z_array = np.logspace(-2, 1, 50)  # z from 0.01 to 10
d_L_std = Planck18.luminosity_distance(z_array)
d_L_custom = my_cosmo.luminosity_distance(z_array)
d_L_phantom = phantom_cosmo.luminosity_distance(z_array)

print(f"Distance difference at z=1:")
print(f"Custom vs Planck18: {(my_cosmo.luminosity_distance(1) - Planck18.luminosity_distance(1)).to(u.Mpc)}")

from astropy.cosmology import z_at_value, Planck18
import astropy.units as u

# Find redshift for specific distance
target_distance = 1000 * u.Mpc
z_target = z_at_value(Planck18.comoving_distance, target_distance)
print(f"Redshift for comoving distance {target_distance}: z = {z_target:.3f}")

# Find redshift for specific age
target_age = 2 * u.Gyr  # 2 billion years after Big Bang
z_young = z_at_value(Planck18.age, target_age)
print(f"Redshift when universe was {target_age}: z = {z_young:.1f}")

# Find redshift for specific lookback time
lookback_target = 10.5 * u.Gyr  # 10.5 Gyr ago
z_lookback = z_at_value(Planck18.lookback_time, lookback_target)
print(f"Redshift for lookback time {lookback_target}: z = {z_lookback:.2f}")

# Calculate angular scales for observations
z_galaxy = 2.0

# Physical scale per angular scale
kpc_per_arcmin = cosmo.kpc_proper_per_arcmin(z_galaxy)
kpc_per_arcsec = kpc_per_arcmin / 60

print(f"At z = {z_galaxy}:")
print(f"  {kpc_per_arcsec:.1f} per arcsecond")
print(f"  {kpc_per_arcmin:.1f} per arcminute")

# Convert observed angular size to physical size
observed_size = 5 * u.arcsec  # 5 arcsecond diameter galaxy
physical_size = observed_size.to(u.kpc, 
                               equivalencies=cosmo.angular_diameter_distance_z1z2(0, z_galaxy))
print(f"5 arcsec corresponds to {physical_size:.1f} at z={z_galaxy}")

# Calculate comoving volumes
z_survey = 3.0

# Total comoving volume out to redshift
total_volume = cosmo.comoving_volume(z_survey)
print(f"Comoving volume out to z={z_survey}: {total_volume.to(u.Gpc**3)}")

# Differential volume element
dV_dz = cosmo.differential_comoving_volume(z_survey)
print(f"Differential volume at z={z_survey}: {dV_dz.to(u.Gpc**3)}")

# Volume of a survey
solid_angle = 10 * u.deg**2  # 10 square degree survey
survey_volume = (solid_angle / (4*np.pi*u.sr)) * total_volume
print(f"Survey volume ({solid_angle}): {survey_volume.to(u.Gpc**3)}")

from astropy.cosmology import Planck18, WMAP9, cosmology_equal

# Compare parameter values
print(f"Planck18 H0: {Planck18.H0}")
print(f"WMAP9 H0: {WMAP9.H0}")
print(f"Planck18 Om0: {Planck18.Om0}")
print(f"WMAP9 Om0: {WMAP9.Om0}")

# Test equality
are_equal = cosmology_equal(Planck18, WMAP9)
print(f"Planck18 == WMAP9: {are_equal}")

# Calculate differences in derived quantities
z_test = 1.0
diff_age = Planck18.age(z_test) - WMAP9.age(z_test)
diff_dL = Planck18.luminosity_distance(z_test) - WMAP9.luminosity_distance(z_test)

print(f"Age difference at z=1: {diff_age.to(u.Myr)}")
print(f"Luminosity distance difference: {diff_dL.to(u.Mpc)}")