Cosmology

cosmology

Cosmology sub-package.

Re-exports the FRW utilities from :mod:jaxlatt.core.cosmology.frw so that existing import sites from jaxlatt.core.cosmology import … keep working unchanged.

FRWUniverse

Bases: Module

FLRW universe parameters in conformal time.

Attributes:

Name	Type	Description
a	Array	Scale factor (dimensionless, typically normalized to \(a(\tau_0) = 1\))
adot	Array	Conformal time derivative \(da/d\tau\) (same units as a/time)
tau	Array	Conformal time coordinate \(\tau\)
M_pl	float	Reduced Planck mass in simulation units (sets gravity strength)

Derived quantities:

• \(H = a'/a\): Conformal Hubble parameter
• \(\rho_{\mathrm{crit}} = 3M_{\mathrm{pl}}^2 H^2\): Critical density

H property

Compute the conformal Hubble parameter.

Returns:

Type	Description
float	Value of \(H = a'/a\) for the current universe state.

rho_crit property

Compute the critical density.

Returns:

Type	Description
float	Critical density computed from 3 * H**2 / (8 * pi * G).

copy()

Create a copy of the universe state.

Returns:

Type	Description
	New FRWUniverse with identical values.

Source code in jaxlatt/core/cosmology/frw.py

def copy(self):
    """Create a copy of the universe state.

    Returns:
        New `FRWUniverse` with identical values.
    """
    return dataclasses.replace(self)

replace(**kwargs)

Create a new universe with selected fields replaced.

Parameters:

Name	Type	Description	Default
**kwargs		Fields to override (a, adot, tau).	{}

Returns:

Type	Description
FRWUniverse	New FRWUniverse with replaced values.

Source code in jaxlatt/core/cosmology/frw.py

def replace(self, **kwargs) -> "FRWUniverse":
    """Create a new universe with selected fields replaced.

    Args:
        **kwargs: Fields to override (a, adot, tau).

    Returns:
        New `FRWUniverse` with replaced values.
    """
    return dataclasses.replace(self, **kwargs)

create_frw_universe(a_initial=1.0, rho_initial=1.0, M_pl=1.0, *, w=1.0 / 3.0)

Initialize a generic FRW universe for a perfect fluid with equation of state \(p = w \rho\).

The first Friedmann equation in conformal time,

\[\left(\frac{a'}{a}\right)^2 = \frac{8\pi G}{3} \rho a^2,\]

fixes the initial \(a'\) from \(\rho\) and \(a\) alone — \(w\) does not enter the initial conditions. Pass the matching pressure function to coupled_evolve_expanding to realize the intended dynamics:

• \(w = 1/3\) (radiation): use make_radiation_pressure
• \(w = 0\) (matter): use make_matter_pressure
• arbitrary: use make_eos_pressure(w)

Parameters:

Name	Type	Description	Default
a_initial	float	Initial scale factor.	1.0
rho_initial	float	Initial energy density.	1.0
M_pl	float	Reduced Planck mass.	1.0
w	float	Equation-of-state parameter \(p/\rho\). Does not affect the returned object; provided for call-site documentation.	1.0 / 3.0

Returns:

Type	Description
FRWUniverse	Initialized FRWUniverse.

Source code in jaxlatt/core/cosmology/frw.py

def create_frw_universe(
    a_initial: float = 1.0,
    rho_initial: float = 1.0,
    M_pl: float = 1.0,
    *,
    w: float = 1.0 / 3.0,
) -> FRWUniverse:
    r"""
    Initialize a generic FRW universe for a perfect fluid with equation of state $p = w \rho$.

    The first Friedmann equation in conformal time,

    $$\left(\frac{a'}{a}\right)^2 = \frac{8\pi G}{3} \rho a^2,$$

    fixes the initial $a'$ from $\rho$ and $a$ alone — $w$ does not enter
    the initial conditions. Pass the matching pressure function to
    ``coupled_evolve_expanding`` to realize the intended dynamics:

    - $w = 1/3$ (radiation): use ``make_radiation_pressure``
    - $w = 0$ (matter): use ``make_matter_pressure``
    - arbitrary: use ``make_eos_pressure(w)``

    Args:
        a_initial: Initial scale factor.
        rho_initial: Initial energy density.
        M_pl: Reduced Planck mass.
        w: Equation-of-state parameter $p/\rho$. Does not affect the returned
            object; provided for call-site documentation.

    Returns:
        Initialized ``FRWUniverse``.
    """
    G = 1.0 / (M_pl * M_pl)
    H_conformal_squared = (8.0 * jnp.pi * G / 3.0) * rho_initial * (a_initial**2)
    adot_initial = a_initial * jnp.sqrt(H_conformal_squared)
    return FRWUniverse(
        a=jnp.asarray(a_initial),
        adot=jnp.asarray(adot_initial),
        tau=jnp.asarray(0.0),
        M_pl=M_pl,
    )

create_matter_universe(a_initial=1.0, rho_initial=1.0, M_pl=1.0)

Initialize FRW universe in matter-dominated era (\(w = 0\)).

See create_frw_universe for full documentation. Pass make_matter_pressure to coupled_evolve_expanding for matter dynamics.

Source code in jaxlatt/core/cosmology/frw.py

def create_matter_universe(
    a_initial: float = 1.0,
    rho_initial: float = 1.0,
    M_pl: float = 1.0,
) -> FRWUniverse:
    """Initialize FRW universe in matter-dominated era ($w = 0$).

    See `create_frw_universe` for full documentation.
    Pass ``make_matter_pressure`` to ``coupled_evolve_expanding`` for matter dynamics.
    """
    return create_frw_universe(a_initial, rho_initial, M_pl, w=0.0)

create_radiation_universe(a_initial=1.0, rho_initial=1.0, M_pl=1.0)

Initialize FRW universe in radiation-dominated era (\(w = 1/3\)).

See create_frw_universe for full documentation. Pass make_radiation_pressure to coupled_evolve_expanding for radiation dynamics.

Source code in jaxlatt/core/cosmology/frw.py

def create_radiation_universe(
    a_initial: float = 1.0,
    rho_initial: float = 1.0,
    M_pl: float = 1.0,
) -> FRWUniverse:
    """Initialize FRW universe in radiation-dominated era ($w = 1/3$).

    See `create_frw_universe` for full documentation.
    Pass ``make_radiation_pressure`` to ``coupled_evolve_expanding`` for radiation dynamics.
    """
    return create_frw_universe(a_initial, rho_initial, M_pl, w=1.0 / 3.0)

friedmann_acceleration(rho, p, M_pl, a)

Compute conformal acceleration \(a''\) from Friedmann equation.

In conformal time \(\tau\) (related to cosmic time \(t\) by \(dt = a(\tau)d\tau\)), the second Friedmann equation is:

\[a'' = \frac{4\pi G}{3} a^3 (\rho - 3p)\]

where:

• \(a'' = d^2a/d\tau^2\)
• \(G = 1/M_{\mathrm{pl}}^2\) (gravitational constant)
• \(\rho\) = energy density
• \(p\) = pressure

Special cases:

• Radiation (\(p = \rho/3\)): \(a'' = 0\) \(\Rightarrow\) \(a(\tau)\) linear in \(\tau\)
• Matter (\(p = 0\)): \(a'' = \frac{4\pi G}{3} a^3 \rho\) \(\Rightarrow\) \(a(\tau) \propto \tau^2\)

References

• Dodelson & Schmidt, "Modern Cosmology" (2nd ed.), Eq. (2.26)
• Mukhanov, "Physical Foundations of Cosmology", Eq. (2.18)
• Baumann, "Cosmology", Eq. (2.35)

Parameters:

Name	Type	Description	Default
rho	float	Total energy density	required
p	float	Total pressure	required
M_pl	float	Reduced Planck mass (default: 1.0 in natural units)	required
a	float	Current scale factor	required

Returns:

Type	Description
Array	Conformal acceleration \(a''\) (second derivative with respect to \(\tau\))

Source code in jaxlatt/core/cosmology/frw.py

@jit
def friedmann_acceleration(rho: float, p: float, M_pl: float, a: float) -> Array:
    """
    Compute conformal acceleration $a''$ from Friedmann equation.

    In conformal time $\\tau$ (related to cosmic time $t$ by $dt = a(\\tau)d\\tau$),
    the second Friedmann equation is:

    $$a'' = \\frac{4\\pi G}{3} a^3 (\\rho - 3p)$$

    where:

    - $a'' = d^2a/d\\tau^2$
    - $G = 1/M_{\\mathrm{pl}}^2$ (gravitational constant)
    - $\\rho$ = energy density
    - $p$ = pressure

    Special cases:

    - Radiation ($p = \\rho/3$): $a'' = 0$ $\\Rightarrow$ $a(\\tau)$ linear in $\\tau$
    - Matter ($p = 0$): $a'' = \\frac{4\\pi G}{3} a^3 \\rho$ $\\Rightarrow$ $a(\\tau) \\propto \\tau^2$

    References:
        - Dodelson & Schmidt, "Modern Cosmology" (2nd ed.), Eq. (2.26)
        - Mukhanov, "Physical Foundations of Cosmology", Eq. (2.18)
        - Baumann, "Cosmology", Eq. (2.35)

    Args:
        rho: Total energy density
        p: Total pressure
        M_pl: Reduced Planck mass (default: 1.0 in natural units)
        a: Current scale factor

    Returns:
        Conformal acceleration $a''$ (second derivative with respect to $\\tau$)
    """
    G = 1.0 / (M_pl * M_pl)
    addot = (4.0 * jnp.pi * G / 3.0) * (rho - 3.0 * p) * (a**3)
    return addot

friedmann_step_leapfrog(universe, rho_func, pressure_func, dt)

Advance scale factor using leapfrog integration (symplectic).

Leapfrog for \(a(\tau)\):

\[a'(\tau + d\tau/2) = a'(\tau) + \frac{d\tau}{2} a''(\tau)\]

\[a(\tau + d\tau) = a(\tau) + d\tau \, a'(\tau + d\tau/2)\]

\[a'(\tau + d\tau) = a'(\tau + d\tau/2) + \frac{d\tau}{2} a''(\tau + d\tau)\]

This preserves symplectic structure if rho_func and pressure_func are Hamiltonian.

Parameters:

Name	Type	Description	Default
universe	FRWUniverse	Current FRWUniverse state	required
rho_func	Callable[[float], float]	Function computing \(\rho(a)\)	required
pressure_func	Callable[[float], float]	Function computing \(p(a)\)	required
dt	float	Conformal timestep \(d\tau\)	required

Returns:

Type	Description
FRWUniverse	Updated FRWUniverse

Source code in jaxlatt/core/cosmology/frw.py

def friedmann_step_leapfrog(
    universe: FRWUniverse,
    rho_func: Callable[[float], float],
    pressure_func: Callable[[float], float],
    dt: float,
) -> FRWUniverse:
    """
    Advance scale factor using leapfrog integration (symplectic).

    Leapfrog for $a(\\tau)$:

    $$a'(\\tau + d\\tau/2) = a'(\\tau) + \\frac{d\\tau}{2} a''(\\tau)$$

    $$a(\\tau + d\\tau) = a(\\tau) + d\\tau \\, a'(\\tau + d\\tau/2)$$

    $$a'(\\tau + d\\tau) = a'(\\tau + d\\tau/2) + \\frac{d\\tau}{2} a''(\\tau + d\\tau)$$

    This preserves symplectic structure if `rho_func` and `pressure_func` are Hamiltonian.

    Args:
        universe: Current FRWUniverse state
        rho_func: Function computing $\\rho(a)$
        pressure_func: Function computing $p(a)$
        dt: Conformal timestep $d\\tau$

    Returns:
        Updated `FRWUniverse`
    """
    a = universe.a
    adot = universe.adot
    tau = universe.tau
    M_pl = universe.M_pl

    # Half-step velocity
    rho = rho_func(a)
    p = pressure_func(a)
    addot = friedmann_acceleration(rho, p, M_pl, a)
    adot_half = adot + 0.5 * dt * addot

    # Full step position
    a_new = a + dt * adot_half

    # Half-step velocity (second half)
    rho_new = rho_func(a_new)
    p_new = pressure_func(a_new)
    addot_new = friedmann_acceleration(rho_new, p_new, M_pl, a_new)
    adot_new = adot_half + 0.5 * dt * addot_new

    tau_new = tau + dt

    return FRWUniverse(a=a_new, adot=adot_new, tau=tau_new, M_pl=M_pl)

friedmann_step_predictor_corrector(universe, rho_func, pressure_func, dt)

Advance scale factor by one timestep using predictor-corrector.

This uses a 2nd-order accurate predictor-corrector scheme:

1. Predictor: Euler step to estimate \(a(\tau + d\tau)\)
2. Corrector: Trapezoidal rule using predicted value

The energy density \(\rho\) and pressure \(p\) are computed from field values via rho_func and pressure_func, which may depend on \(a\) (for rescaled fields).

Parameters:

Name	Type	Description	Default
universe	FRWUniverse	Current FRWUniverse state	required
rho_func	Callable[[float], float]	Function computing total energy density from scale factor. Signature: rho_func(a) -> rho	required
pressure_func	Callable[[float], float]	Function computing total pressure from scale factor. Signature: pressure_func(a) -> p	required
dt	float	Conformal timestep \(d\tau\)	required

Returns:

Type	Description
FRWUniverse	Updated FRWUniverse at \(\tau + d\tau\)

Source code in jaxlatt/core/cosmology/frw.py

def friedmann_step_predictor_corrector(
    universe: FRWUniverse,
    rho_func: Callable[[float], float],
    pressure_func: Callable[[float], float],
    dt: float,
) -> FRWUniverse:
    """
    Advance scale factor by one timestep using predictor-corrector.

    This uses a 2nd-order accurate predictor-corrector scheme:

    1. Predictor: Euler step to estimate $a(\\tau + d\\tau)$
    2. Corrector: Trapezoidal rule using predicted value

    The energy density $\\rho$ and pressure $p$ are computed from field values via
    `rho_func` and `pressure_func`, which may depend on $a$ (for rescaled fields).

    Args:
        universe: Current FRWUniverse state
        rho_func: Function computing total energy density from scale factor.
                  Signature: ``rho_func(a) -> rho``
        pressure_func: Function computing total pressure from scale factor.
                       Signature: ``pressure_func(a) -> p``
        dt: Conformal timestep $d\\tau$

    Returns:
        Updated `FRWUniverse` at $\\tau + d\\tau$
    """

    def rho_p_callable(u: FRWUniverse):
        return rho_func(u.a), pressure_func(u.a)

    return friedmann_step_predictor_corrector_from_state(universe, rho_p_callable, dt)

friedmann_step_predictor_corrector_from_state(universe, rho_p_callable, dt)

Advance scale factor using predictor-corrector with state-dependent ρ, p.

More general than :func:friedmann_step_predictor_corrector: instead of separate rho_func(a) and pressure_func(a) callables, it accepts a single rho_p_callable(universe) → (rho, p). Because the callable receives the whole universe rather than just a, it can close over lattice field state, which is what makes self-consistent field-sourced expansion possible: the fields set ρ and p, which drive a(τ), which in turn evolves the fields.

friedmann_step_predictor_corrector is a thin wrapper around this function, so the two share one integrator and cannot drift apart.

The predictor-corrector scheme is identical to :func:friedmann_step_predictor_corrector; the difference is purely in how ρ and p are evaluated.

Parameters:

Name	Type	Description	Default
universe	FRWUniverse	Current FRWUniverse state.	required
rho_p_callable	Callable[[FRWUniverse], tuple]	Callable (FRWUniverse) → (rho, p).	required
dt	float	Conformal timestep.	required

Returns:

Name	Type	Description
Updated	FRWUniverse	class:FRWUniverse at τ + dt.

Source code in jaxlatt/core/cosmology/frw.py

def friedmann_step_predictor_corrector_from_state(
    universe: FRWUniverse,
    rho_p_callable: Callable[["FRWUniverse"], tuple],
    dt: float,
) -> FRWUniverse:
    """Advance scale factor using predictor-corrector with state-dependent ρ, p.

    More general than :func:`friedmann_step_predictor_corrector`: instead of
    separate ``rho_func(a)`` and ``pressure_func(a)`` callables, it accepts a
    single ``rho_p_callable(universe) → (rho, p)``. Because the callable
    receives the whole universe rather than just ``a``, it can close over
    lattice field state, which is what makes self-consistent field-sourced
    expansion possible: the fields set ρ and p, which drive a(τ), which in turn
    evolves the fields.

    ``friedmann_step_predictor_corrector`` is a thin wrapper around this
    function, so the two share one integrator and cannot drift apart.

    The predictor-corrector scheme is identical to
    :func:`friedmann_step_predictor_corrector`; the difference is purely in
    how ρ and p are evaluated.

    Args:
        universe: Current FRWUniverse state.
        rho_p_callable: Callable ``(FRWUniverse) → (rho, p)``.
        dt: Conformal timestep.

    Returns:
        Updated :class:`FRWUniverse` at τ + dt.
    """
    a = universe.a
    adot = universe.adot
    tau = universe.tau
    M_pl = universe.M_pl

    rho_current, p_current = rho_p_callable(universe)
    addot_current = friedmann_acceleration(rho_current, p_current, M_pl, a)

    a_pred = jnp.maximum(a + dt * adot, 1e-12)
    adot_pred = adot + dt * addot_current

    universe_pred = FRWUniverse(a=a_pred, adot=adot_pred, tau=tau, M_pl=M_pl)
    rho_pred, p_pred = rho_p_callable(universe_pred)
    addot_pred = friedmann_acceleration(rho_pred, p_pred, M_pl, a_pred)

    a_new = jnp.maximum(a + dt * 0.5 * (adot + adot_pred), 1e-12)
    adot_new = adot + dt * 0.5 * (addot_current + addot_pred)
    tau_new = tau + dt
    return FRWUniverse(a=a_new, adot=adot_new, tau=tau_new, M_pl=M_pl)

make_eos_pressure(w, rho_func)

Create pressure function for equation of state \(p = w\rho\).

Parameters:

Name	Type	Description	Default
w	float	Equation of state parameter	required
rho_func	Callable[[float], float]	Energy density function \(\rho(a)\)	required

Returns:

Type	Description
	Function \(p(a) = w\rho(a)\)

Source code in jaxlatt/core/cosmology/frw.py

def make_eos_pressure(w: float, rho_func: Callable[[float], float]):
    """
    Create pressure function for equation of state $p = w\\rho$.

    Args:
        w: Equation of state parameter
        rho_func: Energy density function $\\rho(a)$

    Returns:
        Function $p(a) = w\\rho(a)$
    """

    def pressure_func(a: float) -> float:
        return w * rho_func(a)

    return pressure_func

make_matter_pressure(a=None)

Create matter pressure function \(p = 0\).

Parameters:

Name	Type	Description	Default
a	float	Scale factor (unused, for API consistency)	None

Returns:

Type	Description
	Function \(p(a) = 0\) for matter

Source code in jaxlatt/core/cosmology/frw.py

def make_matter_pressure(a: float = None):
    """
    Create matter pressure function $p = 0$.

    Args:
        a: Scale factor (unused, for API consistency)

    Returns:
        Function $p(a) = 0$ for matter
    """

    def pressure_func(a):
        return 0.0 * jnp.asarray(a)

    return pressure_func

make_radiation_pressure(rho_initial, a_initial)

Create radiation pressure function \(p = \rho/3\) with \(\rho \propto a^{-4}\).

Parameters:

Name	Type	Description	Default
rho_initial	float	Initial energy density	required
a_initial	float	Initial scale factor	required

Returns:

Type	Description
	Function \(p(a)\) for radiation

Source code in jaxlatt/core/cosmology/frw.py

def make_radiation_pressure(rho_initial: float, a_initial: float):
    """
    Create radiation pressure function $p = \\rho/3$ with $\\rho \\propto a^{-4}$.

    Args:
        rho_initial: Initial energy density
        a_initial: Initial scale factor

    Returns:
        Function $p(a)$ for radiation
    """

    def pressure_func(a: float) -> float:
        rho = rho_initial * (a_initial / a) ** 4
        return rho / 3.0

    return pressure_func

make_radiation_rho(rho_initial, a_initial)

Create a radiation-era density function with \(\rho(a) \propto a^{-4}\).

Parameters:

Name	Type	Description	Default
rho_initial	float	Reference density at a_initial.	required
a_initial	float	Reference scale factor.	required

Returns:

Type	Description
	Callable rho_func(a) that evaluates radiation density at scale factor a.

Source code in jaxlatt/core/cosmology/frw.py

def make_radiation_rho(rho_initial: float, a_initial: float):
    r"""Create a radiation-era density function with $\rho(a) \propto a^{-4}$.

    Args:
        rho_initial: Reference density at `a_initial`.
        a_initial: Reference scale factor.

    Returns:
        Callable `rho_func(a)` that evaluates radiation density at scale factor `a`.
    """

    def rho_func(a: float) -> float:
        return rho_initial * (a_initial / a) ** 4

    return rho_func

matter_scaling_check(a_initial, a_final, rho_initial, rho_final)

Check if energy density scales as matter (\(\rho \propto a^{-3}\)).

Parameters:

Name	Type	Description	Default
a_initial	float	Initial scale factor.	required
a_final	float	Final scale factor.	required
rho_initial	float	Initial energy density.	required
rho_final	float	Final energy density.	required

Returns:

Type	Description
float	Tuple (scaling_exponent, relative_error) where scaling_exponent is the
float	measured exponent \(n\) in \(\rho \propto a^{-n}\) and relative_error is
tuple[float, float]	the deviation from \(n = 3\).

Source code in jaxlatt/core/cosmology/frw.py

def matter_scaling_check(
    a_initial: float,
    a_final: float,
    rho_initial: float,
    rho_final: float,
) -> tuple[float, float]:
    """
    Check if energy density scales as matter ($\\rho \\propto a^{-3}$).

    Args:
        a_initial: Initial scale factor.
        a_final: Final scale factor.
        rho_initial: Initial energy density.
        rho_final: Final energy density.

    Returns:
        Tuple `(scaling_exponent, relative_error)` where `scaling_exponent` is the
        measured exponent $n$ in $\\rho \\propto a^{-n}$ and `relative_error` is
        the deviation from $n = 3$.
    """
    return _scaling_check(a_initial, a_final, rho_initial, rho_final, 3.0)

radiation_scaling_check(a_initial, a_final, rho_initial, rho_final)

Check if energy density scales as radiation (\(\rho \propto a^{-4}\)).

For pure radiation:

\[\rho(a) = \rho_0 \left(\frac{a_0}{a}\right)^4\]

Parameters:

Name	Type	Description	Default
a_initial	float	Initial scale factor.	required
a_final	float	Final scale factor.	required
rho_initial	float	Initial energy density.	required
rho_final	float	Final energy density.	required

Returns:

Type	Description
float	Tuple (scaling_exponent, relative_error) where scaling_exponent is the
float	measured exponent \(n\) in \(\rho \propto a^{-n}\) and relative_error is
tuple[float, float]	the deviation from \(n = 4\).

Source code in jaxlatt/core/cosmology/frw.py

def radiation_scaling_check(
    a_initial: float,
    a_final: float,
    rho_initial: float,
    rho_final: float,
) -> tuple[float, float]:
    """
    Check if energy density scales as radiation ($\\rho \\propto a^{-4}$).

    For pure radiation:

    $$\\rho(a) = \\rho_0 \\left(\\frac{a_0}{a}\\right)^4$$

    Args:
        a_initial: Initial scale factor.
        a_final: Final scale factor.
        rho_initial: Initial energy density.
        rho_final: Final energy density.

    Returns:
        Tuple `(scaling_exponent, relative_error)` where `scaling_exponent` is the
        measured exponent $n$ in $\\rho \\propto a^{-n}$ and `relative_error` is
        the deviation from $n = 4$.
    """
    return _scaling_check(a_initial, a_final, rho_initial, rho_final, 4.0)