Skip to content

Friedmann expansion

The scale factor of a homogeneous universe of matter and dark energy, after Friedmann and Lemaître.

  • Name: friedmann_expansion
  • Field: physics/cosmology
  • Formalism: ordinary differential equations
  • References: Friedmann 1922; Lemaître 1927; Satz 2022

Factory

@register_model(
    field='physics/cosmology',
    tags=('cosmology', 'ode', 'second order', 'gravity', 'dark energy'),
    references=('Friedmann 1922', 'Lemaître 1927', 'Satz 2022'),
    defaults={'until': 2.0, 'dt': 0.02},
)
def friedmann_expansion(
    a: float = 1.0,
    *,
    H0: float = 1.0,
    omega_matter: float = 0.3,
    omega_lambda: float = 0.7,
) -> System:
    """Return the scale factor of a homogeneous universe of matter and dark energy, after Friedmann and Lemaître.

    The acceleration equation `a² a'' = -4 pi G M / 3 + Lambda a³ / 3` is written with the density parameters of
    matter `Ω_m = 8 pi G M / (3 H0²)` and of the cosmological constant `Ω_Λ = Lambda / (3 H0²)` in units of the Hubble
    constant `H0` at `a = 1`, so `a² a'' = -Ω_m H0² / 2 + Ω_Λ H0² a³`. The initial expansion rate follows from the
    Friedmann equation `a'² = H0² (Ω_m / a + Ω_Λ a² + 1 - Ω_m - Ω_Λ)`, which stays a first integral of the motion.
    Matter decelerates the expansion and dark energy accelerates it once `a` exceeds `(Ω_m / (2 Ω_Λ))^(1/3)`.

    Args:
        a:
            The initial scale factor.

        H0:
            The Hubble constant, the expansion rate at `a = 1`.

        omega_matter:
            The density parameter of matter.

        omega_lambda:
            The density parameter of the cosmological constant.

    Returns:
        system:
            The model, with the scale factor `universe.a` of order 2 and its rate `universe.a_d1`.

    Raises:
        ValidationError:
            The universe does not expand at the initial scale factor.

    Examples:
        >>> from skcomplex.models.physics.cosmology import friedmann_expansion
        >>> friedmann_expansion(omega_matter=1.0, omega_lambda=0.0).elements[0].a_d1
        1.0
    """
    square = H0**2 * (omega_matter / a + omega_lambda * a**2 + 1 - omega_matter - omega_lambda)
    if square < 0:
        error_msg = (
            f'The Friedmann equation gives a negative squared expansion rate {square} at a = {a}. Choose a, '
            'omega_matter, and omega_lambda for which it is positive.'
        )
        raise ValidationError(error_msg)
    universe = ScaleFactorElement(
        'universe',
        a=a,
        a_d1=math.sqrt(square),
        H0=H0,
        omega_matter=omega_matter,
        omega_lambda=omega_lambda,
    )
    acceleration = GradientInteraction(
        'acceleration',
        target='universe',
        states='a',
        potential=_get_expansion_potential,
    )
    return System(elements=[universe], interactions=[acceleration])

Components

class ScaleFactorElement(Element):
    """Homogeneous universe whose size is the scale factor `a`, a state of order 2 with the expansion rate `a_d1`.

    Args:
        label:
            The label of the universe.

        visual:
            How the universe is drawn, or `None` for none.

        a:
            The initial scale factor.

        a_d1:
            The initial rate of change of the scale factor.

        H0:
            The Hubble constant.

        omega_matter:
            The density parameter of matter.

        omega_lambda:
            The density parameter of dark energy.

    Examples:
        >>> from skcomplex.models.physics.cosmology import ScaleFactorElement
        >>> ScaleFactorElement('universe', H0=2.0).H0
        2.0
    """

    def __init__(
        self,
        label: str,
        visual: VisualSpec | None = None,
        *,
        a: Annotated[float, StateSpec(domain=RealSet.open(0, oo), order=2)] = 1.0,
        a_d1: float = 0.0,
        H0: Annotated[float, ParameterSpec(domain=RealSet.open(0, oo))] = 1.0,
        omega_matter: Annotated[float, ParameterSpec(domain=RealSet.closed_open(0, oo))] = 0.3,
        omega_lambda: Annotated[float, ParameterSpec()] = 0.7,
    ) -> None:
        super().__init__(label, visual=visual)
        self.a = a
        self.a_d1 = a_d1
        self.H0 = H0
        self.omega_matter = omega_matter
        self.omega_lambda = omega_lambda

class GradientInteraction(Interaction):
    """Motion of states down the gradient of a potential `U(ns)`, as a gradient flow or as Newtonian motion.

    Every state `x` the interaction moves gets the rate `-dU/dx / d` on its highest time derivative, with `d` the
    divisor of its member: a state or parameter of the member, an expression of them such as a moment of inertia
    `m l^2`, or one divisor per state, such as the masses of two coordinates of one element. On a state of order 1
    this is the gradient flow `x' = -dU/dx / d`, with `d` a friction, and on a state of order 2 it is Newton's law
    `x'' = -dU/dx / d`, with `d` an inertia. A vector state gets the gradient of every component, and a tensor state
    the matrix gradient, the projection of the derivative onto the matrices of its symmetry.

    On a field the potential is the density `f` of a functional `F = integral f`, which may read the field, its
    derivatives `partial(x, k)` or `gradient(x)`, and its `laplacian(x)`. A variational interaction moves the field by
    the functional derivative `dF/dx = df/dx - sum_k partial(df/dx_k, k) + laplacian(df/d laplacian(x))`, as in the
    Allen-Cahn equation, and a conserved interaction gives the rate `laplacian(dF/dx) / d` instead, as in the
    Cahn-Hilliard equation, so the total of the field stays constant.

    The interaction moves the states of every role that binds one element or group, `target` by default. A subclass
    may bind several such roles instead, such as both members of a pair potential, and every one of them moves. The
    potential function takes the namespace of the interaction, with the namespaces of the roles and of the `sources`,
    the time `t`, and the symbols of the parameters a subclass declares.

    Args:
        label:
            The label of the interaction, unique among the parts of its system.

        target:
            The path of the element or group whose states move.

        states:
            The name of the state that moves, or the names of several states.

        potential:
            A function of the namespace that returns the potential, or the density of the functional of a field.

        sources:
            The paths of the elements or groups the potential reads but does not move.

        divisor:
            The name of the state or parameter of every moving member that divides its rate, a function of the
            namespace of the member that returns the divisor, a mapping from the name of every moving state to its
            own divisor, or `None`.

        variational:
            Whether the potential is the density of a functional of a field, moved by its functional derivative.

        conserved:
            Whether the rate is the Laplacian of the derivative of the potential, which conserves the total of a field.

    Raises:
        ValidationError:
            A state or the divisor is not a name, or the potential is not callable.

    Examples:
        >>> from skcomplex.core.elements import CompartmentElement
        >>> from skcomplex.core.interactions import GradientInteraction
        >>> from skcomplex.core.system import System
        >>> from skcomplex.engine.runtime import Simulation
        >>> well = GradientInteraction('well', target='walker', states='n', potential=lambda s: s.target.n**2 / 2)
        >>> system = System(elements=[CompartmentElement('walker', n=1.0)], interactions=[well])
        >>> trajectory = Simulation(system, dt=0.01).check().prepare().run(until=1.0).trajectory_
        >>> round(float(trajectory.values('walker.n')[-1]), 4)
        0.3679
    """

    def __init__(
        self,
        label: str,
        *,
        target: Annotated[str, RoleSpec()],
        states: str | Sequence[str] = 'x',
        potential: Callable[[SimpleNamespace], object] | None = None,
        sources: Annotated[Sequence[str], RoleSpec(variadic=True)] = (),
        divisor: (
            str | Callable[[SimpleNamespace], object] | Mapping[str, str | Callable[[SimpleNamespace], object]] | None
        ) = None,
        variational: bool = False,
        conserved: bool = False,
    ) -> None:
        super().__init__(label)
        self.target = target
        self.states = states
        self.potential = potential
        self.sources = sources
        self.divisor = divisor
        self.variational = variational
        self.conserved = conserved

    def _get_names(self) -> tuple[str, ...]:
        """Return the names of the states the interaction moves."""
        return (self.states,) if isinstance(self.states, str) else tuple(self.states)

    def _check(self) -> None:
        """Raise an error for a state or a divisor that is not a name, or a potential that is not callable.

        Raises:
            ValidationError:
                The type declares states, a state is not a name, a divisor is neither a name nor a function, or the
                potential is not callable.
        """
        super()._check()
        names = self._get_names()
        if not names or not all(isinstance(name, str) and name.isidentifier() for name in names):
            error_msg = f'{type(self).__name__} {self.label!r} moves the states {self.states!r}. Give names of states.'
            raise ValidationError(error_msg)
        divisors = dict(self.divisor) if isinstance(self.divisor, Mapping) else dict.fromkeys(names, self.divisor)
        valid = all(
            item is None or callable(item) or (isinstance(item, str) and item.isidentifier())
            for item in divisors.values()
        )
        if not valid or not set(divisors) <= set(names):
            error_msg = (
                f'The divisor of {self.label!r} is {self.divisor!r}. Give the name of a state or parameter, a function '
                'of the namespace of the member, or a mapping from the moving states to them.'
            )
            raise ValidationError(error_msg)
        if self.potential is not None and not callable(self.potential):
            error_msg = f'The potential of {self.label!r} is {self.potential!r}. Give a function of the namespace.'
            raise ValidationError(error_msg)

    def _get_potential(self, namespace: SimpleNamespace) -> object:
        """Return the potential, the hook a subclass overrides to derive it from its parameters."""
        if self.potential is None:
            error_msg = f'{type(self).__name__} {self.label!r} has no potential. Give a function of the namespace.'
            raise ValidationError(error_msg)
        return self.potential(namespace)

    def _get_divisor(self, member: SimpleNamespace, name: str) -> object:
        """Return the divisor of a moving state of a member, or `None` for none."""
        divisor = self.divisor.get(name) if isinstance(self.divisor, Mapping) else self.divisor
        if divisor is None:
            return None
        return getattr(member, divisor) if isinstance(divisor, str) else divisor(member)

    def _get_rate_state(self, member: SimpleNamespace, name: str) -> object:
        """Return the state whose highest derivative gets the rate of a moving state, the moving state itself."""
        return getattr(member, name)

    def _walk(self, expression: Expression) -> Iterator[Expression]:
        """Yield an expression and every subexpression of it."""
        yield expression
        for operand in expression.operands():
            yield from self._walk(operand)

    def _get_flux_divergence(
        self,
        fluxes: dict[int, Expression],
        derivatives: dict[int, Expression],
        field: Expression,
    ) -> Expression:
        """Return `sum_k partial(flux_k, k)`, or `c laplacian(x)` when every flux is `c x_k` along every axis read.

        The Laplacian keeps the compact stencil of the field, which nested first derivatives would widen.
        """
        coefficients = [SR(fluxes[axis] / derivatives[axis]) for axis in sorted(fluxes)]
        coefficient = coefficients[0]
        uniform = sorted(fluxes) == list(range(len(fluxes))) and all(
            item.is_trivially_equal(coefficient) for item in coefficients
        )
        constant = not coefficient.has(field) and not any(coefficient.has(item) for item in derivatives.values())
        if uniform and constant:
            return coefficient * laplacian(field)
        return sum((partial(flux, SR(axis)) for axis, flux in fluxes.items()), SR(0))

    def _get_variation(self, density: Expression, field: Expression) -> Expression:
        """Return the functional derivative of the integral of a density with respect to a field."""
        derivatives: dict[int, Expression] = {}
        curvature = None
        for item in self._walk(density):
            kind = item.operator()
            if kind == partial and item.operands()[0].is_trivially_equal(field):
                derivatives[int(item.operands()[1])] = item
            elif kind == laplacian and item.operands()[0].is_trivially_equal(field):
                curvature = item
        slots = {item: SR.var(f'_slot_{axis}') for axis, item in derivatives.items()}
        if curvature is not None:
            slots[curvature] = SR.var('_slot_laplacian')
        flat = density.subs(slots)
        back = {symbol: item for item, symbol in slots.items()}
        variation = flat.diff(field).subs(back)
        fluxes = {axis: flat.diff(slots[item]).subs(back) for axis, item in derivatives.items()}
        if fluxes:
            variation -= self._get_flux_divergence(fluxes, derivatives, field)
        if curvature is not None:
            variation += self._get_laplacian(flat.diff(slots[curvature]).subs(back), field)
        return variation

    def _get_laplacian(self, expression: Expression, field: Expression) -> Expression:
        """Return the Laplacian of an expression of a field, term by term, with the factors free of the field outside.

        The linear terms of the field, such as `laplacian(x)`, stay apart, so a spectral discretization finds them.
        """
        expanded = SR(expression).expand()
        terms = expanded.operands() if expanded.operator() is add_vararg else [expanded]
        total = SR(0)
        for term in terms:
            factors = term.operands() if term.operator() is mul_vararg else [term]
            constant = prod((factor for factor in factors if not factor.has(field)), start=SR(1))
            varying = prod((factor for factor in factors if factor.has(field)), start=SR(1))
            if varying.has(field):
                total += constant * laplacian(varying)
        return total

    def _get_matrix_derivative(self, potential: Expression, state: Matrix) -> Matrix:
        """Return the matrix gradient of a potential at a tensor state, written in the matrices of its symmetry.

        The derivative with respect to the independent components is the product of the matrices of the components, the
        metric of the state, and the coefficients of the gradient, so the metric gives the coefficients.
        """
        if self.conserved:
            error_msg = f'{self.label!r} conserves the tensor state {state}. Give a scalar state.'
            raise ValidationError(error_msg)
        components = sorted(state.variables(), key=str)
        units = [state.derivative(component) for component in components]
        metric = matrix(SR, [[(first.transpose() * second).trace() for second in units] for first in units])
        density = distribute_operators(potential)
        derivatives = [
            self._get_variation(density, component) if self.variational else density.diff(component)
            for component in components
        ]
        coefficients = metric.inverse() * vector(SR, derivatives)
        total = matrix(SR, state.nrows(), state.ncols())
        for coefficient, unit in zip(coefficients, units, strict=True):
            total += coefficient * unit
        return total

    def _get_derivative(self, potential: Expression, state: Target | Matrix) -> Target | Matrix:
        """Return the derivative of the potential with respect to a state, every component of a vector state."""
        if isinstance(state, Matrix):
            return self._get_matrix_derivative(potential, state)
        if isinstance(state, FreeModuleElement):
            if self.variational or self.conserved:
                error_msg = f'{self.label!r} moves the vector state {state} of a field. Give a scalar state.'
                raise ValidationError(error_msg)
            return vector(SR, [potential.diff(component) for component in state])
        if self.variational:
            return self._get_variation(potential, state)
        return potential.diff(state)

    def _get_members(self, namespace: SimpleNamespace) -> list[SimpleNamespace]:
        """Return the namespaces of the members whose states move, those of every role that binds one path."""
        return [
            getattr(namespace, name)
            for name, spec in self._get_specs().items()
            if isinstance(spec, RoleSpec) and not spec.variadic and not spec.noise
        ]

    def _get_laws(self, namespace: SimpleNamespace) -> list[BasePart]:
        """Return the rate of every moving state, minus the derivative of the potential over the divisor.

        Args:
            namespace:
                The namespace of the symbols the laws can use, with the namespaces of the roles.

        Returns:
            laws:
                The rate laws of the moving states.
        """
        potential = SR(self._get_potential(namespace))
        laws: list[BasePart] = []
        for member in self._get_members(namespace):
            for name in self._get_names():
                state = getattr(member, name)
                derivative = self._get_derivative(potential, state)
                rate = self._get_laplacian(derivative, state) if self.conserved else -derivative
                divisor = self._get_divisor(member, name)
                if divisor is not None:
                    rate = rate / divisor
                laws.append(RateLaw(self._get_rate_state(member, name), rate))
        return laws

Equations

universe.a: diff(a_universe(t), t, t) == H0_universe^2*omega_lambda_universe*a_universe(t) - 1/2*H0_universe^2*omega_matter_universe/a_universe(t)^2