physicskit 0.1.0__py3-none-any.whl

physicskit/astro/stellar_structure.py

@@ -0,0 +1,198 @@
+r"""Polytropic stellar models via the Lane-Emden equation.
+
+Uses **gravitational units with** :math:`G=1` by default -- every function
+that needs it accepts ``G`` as a keyword argument, the same convention
+:mod:`physicskit.relativity` uses for its own geometrized units.
+
+- :func:`lane_emden` -- the dimensionless hydrostatic-equilibrium profile
+  of a self-gravitating polytropic sphere.
+- :class:`PolytropicStar` -- a physical star (radius, mass) built from one
+  Lane-Emden solution.
+- :func:`chandrasekhar_mass` -- the white-dwarf mass limit.
+- :func:`main_sequence_luminosity` -- the empirical mass-luminosity relation.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from scipy.integrate import solve_ivp
+
+__all__ = ["lane_emden", "PolytropicStar", "chandrasekhar_mass", "main_sequence_luminosity"]
+
+
+def lane_emden(n, xi_max=10.0, n_points=2000):
+    r"""Integrate the Lane-Emden equation for a polytrope of index ``n``.
+
+    .. math::
+
+        \frac{1}{\xi^2}\frac{d}{d\xi}\left(\xi^2\frac{d\theta}{d\xi}\right)
+        + \theta^n = 0, \qquad \theta(0)=1,\ \theta'(0)=0
+
+    the dimensionless hydrostatic-equilibrium profile of a self-gravitating
+    sphere with equation of state :math:`P=K\rho^{1+1/n}`.
+
+    Parameters
+    ----------
+    n : float
+        Polytropic index.
+    xi_max : float, default=10.0
+        Upper integration bound, used only as a fallback for indices
+        (e.g. ``n=5``) where :math:`\theta` never reaches zero.
+    n_points : int, default=2000
+        Number of points to sample the returned arrays at.
+
+    Returns
+    -------
+    xi : ndarray
+        Dimensionless radius, from just above 0 to the surface (the first
+        zero of :math:`\theta`) or to ``xi_max``.
+    theta : ndarray
+        The Lane-Emden function :math:`\theta(\xi)`.
+
+    Examples
+    --------
+    >>> xi, theta = lane_emden(0.0)
+    >>> round(float(xi[-1]), 3)  # exact analytic surface is sqrt(6)
+    2.449
+    """
+    xi0 = 1e-6
+    theta0 = 1.0 - xi0**2 / 6.0
+    dtheta0 = -xi0 / 3.0
+
+    def rhs(xi, y):
+        theta, dtheta = y
+        theta_pos = max(theta, 0.0)
+        d2theta = -(theta_pos**n) - (2.0 / xi) * dtheta
+        return [dtheta, d2theta]
+
+    def surface_event(xi, y):
+        return y[0]
+
+    surface_event.terminal = True
+    surface_event.direction = -1
+
+    sol = solve_ivp(
+        rhs,
+        [xi0, xi_max],
+        [theta0, dtheta0],
+        events=surface_event,
+        dense_output=True,
+        rtol=1e-9,
+        atol=1e-11,
+        max_step=xi_max / 2000,
+    )
+    xi_surface = sol.t[-1]
+    xi = np.linspace(xi0, xi_surface, n_points)
+    theta = sol.sol(xi)[0]
+    theta = np.clip(theta, 0.0, None)
+    return xi, theta
+
+
+class PolytropicStar:
+    r"""A physical polytropic star built from one Lane-Emden solution.
+
+    Parameters
+    ----------
+    n : float
+        Polytropic index, :math:`P=K\rho^{1+1/n}`.
+    K : float
+        Polytropic constant.
+    rho_c : float
+        Central density.
+    G : float, default=1.0
+        Gravitational constant.
+
+    Examples
+    --------
+    >>> star = PolytropicStar(n=0.0, K=1.0, rho_c=1.0)
+    >>> round(star.mass / (4 / 3 * np.pi * star.radius**3 * star.rho_c), 3)
+    1.0
+    """
+
+    def __init__(self, n, K, rho_c, G=1.0):
+        self.n = n
+        self.K = K
+        self.rho_c = rho_c
+        self.G = G
+        self._xi, self._theta = lane_emden(n)
+
+    @property
+    def xi1(self):
+        """The dimensionless surface radius, the last point of the Lane-Emden solution."""
+        return float(self._xi[-1])
+
+    @property
+    def alpha(self):
+        r"""Physical length scale, :math:`\alpha=\sqrt{(n+1)K\rho_c^{1/n-1}/(4\pi G)}`."""
+        if self.n == 0:
+            exponent = -1.0
+        else:
+            exponent = 1.0 / self.n - 1.0
+        return float(np.sqrt((self.n + 1) * self.K * self.rho_c**exponent / (4.0 * np.pi * self.G)))
+
+    @property
+    def radius(self):
+        r"""Physical stellar radius, :math:`R=\alpha\xi_1`."""
+        return self.alpha * self.xi1
+
+    @property
+    def mass(self):
+        r"""Physical stellar mass, :math:`M=4\pi\alpha^3\rho_c\left[-\xi_1^2\theta'(\xi_1)\right]`."""
+        dtheta_surface = (self._theta[-1] - self._theta[-2]) / (self._xi[-1] - self._xi[-2])
+        return float(4.0 * np.pi * self.alpha**3 * self.rho_c * (-(self.xi1**2) * dtheta_surface))
+
+
+def chandrasekhar_mass(mu_e=2.0):
+    r"""The Chandrasekhar mass limit, in solar masses.
+
+    .. math::
+
+        M_{\rm Ch} \approx \frac{5.83}{\mu_e^2}\ M_\odot
+
+    (the standard numerical prefactor from the n=3 polytrope /
+    relativistic-electron-degeneracy treatment; see e.g. Kippenhahn,
+    Weigert & Weiss, *Stellar Structure and Evolution*).
+
+    Parameters
+    ----------
+    mu_e : float, default=2.0
+        Mean molecular weight per electron (2.0 for carbon/oxygen white
+        dwarfs).
+
+    Returns
+    -------
+    float
+        Mass limit, in solar masses.
+
+    Examples
+    --------
+    >>> M = chandrasekhar_mass(2.0)
+    >>> 1.2 < M < 1.5
+    True
+    """
+    return 5.83 / mu_e**2
+
+
+def main_sequence_luminosity(mass_solar):
+    r"""Empirical main-sequence mass-luminosity relation, :math:`L/L_\odot\approx(M/M_\odot)^{3.5}`.
+
+    Valid roughly for main-sequence stars in the 0.5-10 solar mass range;
+    this is a pure empirical power law, not derived from first-principles
+    stellar structure.
+
+    Parameters
+    ----------
+    mass_solar : float
+        Mass, in solar masses.
+
+    Returns
+    -------
+    float
+        Luminosity, in solar luminosities.
+
+    Examples
+    --------
+    >>> round(main_sequence_luminosity(1.0), 6)
+    1.0
+    """
+    return mass_solar**3.5

physicskit/astro/tests/conftest.py

@@ -0,0 +1,7 @@
+"""Test-session setup: force Matplotlib's non-interactive Agg backend
+before anything imports pyplot, so the suite (and CI, which has no
+display) can exercise the plotting/animation code headlessly."""
+
+import matplotlib
+
+matplotlib.use("Agg")

physicskit/astro/tests/test_cosmic_web.py

@@ -0,0 +1,159 @@
+import numpy as np
+import pytest
+
+from physicskit.astro.cosmic_web import (
+    first_caustic_time,
+    lagrangian_grid,
+    random_displacement_potential,
+    zeldovich_displacement,
+    zeldovich_hessian_eigenvalues,
+    zeldovich_position,
+)
+
+# For a single plane-wave potential Phi(q) = A*cos(kx*x + ky*y + phase), the
+# Hessian is H = -A*cos(theta)*[[kx^2, kx*ky], [kx*ky, ky^2]] at every q,
+# i.e. a rank-1 (outer-product) matrix scaled by c = -A*cos(theta). Its two
+# eigenvalues are exactly {0, c*|k|^2}, so the *largest* eigenvalue at any
+# given q is max(0, c*|k|^2) -- and this is maximized over q (giving the
+# analytic first-caustic time 1/(A*|k|^2)) exactly where cos(theta) = -1,
+# i.e. where kx*x + ky*y + phase = pi (mod 2*pi).
+
+
+def test_single_mode_hessian_eigenvalues_exact_at_compression_maximum():
+    """At the specific q where the mode's phase argument equals pi (its
+    point of maximal compression), the Hessian eigenvalues are exactly
+    (0, A*|k|^2) -- a fully deterministic, hand-derivable check."""
+    kx, ky, A = 2 * np.pi, 0.0, 1.5
+    k_vectors = np.array([[kx, ky]])
+    amplitudes = np.array([A])
+    phases = np.array([np.pi])  # theta = kx*0 + ky*0 + pi = pi at q=(0,0)
+
+    q_star = np.array([[0.0, 0.0]])
+    eigs = zeldovich_hessian_eigenvalues(q_star, k_vectors, amplitudes, phases)
+
+    k_sq = kx**2 + ky**2
+    assert eigs.shape == (1, 2)
+    assert eigs[0, 0] == pytest.approx(0.0, abs=1e-12)
+    assert eigs[0, 1] == pytest.approx(A * k_sq, rel=1e-10)
+
+
+def test_first_caustic_time_matches_single_mode_analytic_value_exactly():
+    """first_caustic_time at that same single deterministic point should
+    equal the closed-form D_collapse = 1/(A*|k|^2) exactly."""
+    kx, ky, A = 2 * np.pi, 0.0, 1.5
+    k_vectors = np.array([[kx, ky]])
+    amplitudes = np.array([A])
+    phases = np.array([np.pi])
+
+    q_star = np.array([[0.0, 0.0]])
+    D_collapse = first_caustic_time(q_star, k_vectors, amplitudes, phases)
+    analytic = 1.0 / (A * (kx**2 + ky**2))
+    assert D_collapse == pytest.approx(analytic, rel=1e-10)
+
+
+def test_first_caustic_time_converges_with_grid_resolution():
+    """On a *generic* grid (the compression maximum does not sit exactly on
+    a grid point), first_caustic_time only finds the nearest sampled point
+    to the true maximum, so it overestimates D_collapse -- but the error
+    should shrink as the grid is refined and the sampling gets closer to
+    the true maximizing point."""
+    L = 1.0
+    kx, ky, A = 2 * np.pi / L, 0.0, 1.0
+    k_vectors = np.array([[kx, ky]])
+    amplitudes = np.array([A])
+    phases = np.array([0.7])  # arbitrary offset, not aligned with any grid spacing
+    analytic = 1.0 / (A * (kx**2 + ky**2))
+
+    q_coarse = lagrangian_grid(20, L)
+    q_fine = lagrangian_grid(400, L)
+
+    D_coarse = first_caustic_time(q_coarse, k_vectors, amplitudes, phases)
+    D_fine = first_caustic_time(q_fine, k_vectors, amplitudes, phases)
+
+    error_coarse = D_coarse - analytic
+    error_fine = D_fine - analytic
+
+    # The grid-based estimate can only ever undershoot the true max
+    # eigenvalue, so it can only ever overshoot D_collapse.
+    assert error_coarse >= 0.0
+    assert error_fine >= 0.0
+    assert error_fine < error_coarse
+    assert error_fine < 1e-5
+
+
+def test_zeldovich_position_at_zero_growth_returns_lagrangian_position():
+    k_vectors, amplitudes, phases = random_displacement_potential(6, k_min=1.0, k_max=4.0, amplitude_scale=0.2, seed=3)
+    q = lagrangian_grid(12, 2.0)
+    x0 = zeldovich_position(q, 0.0, k_vectors, amplitudes, phases)
+    assert x0 == pytest.approx(q, abs=1e-12)
+
+
+def test_zeldovich_displacement_matches_finite_difference_gradient():
+    """Sanity check independent of the eigenvalue-formula correctness:
+    the analytic gradient should match a numerical finite-difference
+    gradient of Phi evaluated by hand from the same plane-wave sum."""
+    rng = np.random.default_rng(42)
+    n_modes = 4
+    k_vectors = rng.uniform(-3.0, 3.0, size=(n_modes, 2))
+    amplitudes = rng.uniform(0.1, 1.0, size=n_modes)
+    phases = rng.uniform(0.0, 2 * np.pi, size=n_modes)
+
+    def phi(point):
+        theta = k_vectors @ point + phases
+        return float(np.sum(amplitudes * np.cos(theta)))
+
+    q_points = rng.uniform(-2.0, 2.0, size=(5, 2))
+    analytic_grad = zeldovich_displacement(q_points, k_vectors, amplitudes, phases)
+
+    h = 1e-6
+    for i, point in enumerate(q_points):
+        numeric_grad = np.zeros(2)
+        for d in range(2):
+            step = np.zeros(2)
+            step[d] = h
+            numeric_grad[d] = (phi(point + step) - phi(point - step)) / (2 * h)
+        assert analytic_grad[i] == pytest.approx(numeric_grad, abs=1e-5)
+
+
+def test_shapes_and_no_nans():
+    k_vectors, amplitudes, phases = random_displacement_potential(5, k_min=1.0, k_max=3.0, amplitude_scale=0.1, seed=7)
+    q = lagrangian_grid(10, 2.0)
+    n = q.shape[0]
+
+    assert q.shape == (100, 2)
+    assert k_vectors.shape == (5, 2)
+    assert amplitudes.shape == (5,)
+    assert phases.shape == (5,)
+
+    grad = zeldovich_displacement(q, k_vectors, amplitudes, phases)
+    x = zeldovich_position(q, 0.5, k_vectors, amplitudes, phases)
+    eigs = zeldovich_hessian_eigenvalues(q, k_vectors, amplitudes, phases)
+    D_collapse = first_caustic_time(q, k_vectors, amplitudes, phases)
+
+    assert grad.shape == (n, 2)
+    assert x.shape == (n, 2)
+    assert eigs.shape == (n, 2)
+    assert np.all(eigs[:, 0] <= eigs[:, 1])
+    assert isinstance(D_collapse, float)
+    assert D_collapse > 0.0
+
+    assert not np.any(np.isnan(grad))
+    assert not np.any(np.isnan(x))
+    assert not np.any(np.isnan(eigs))
+    assert not np.isnan(D_collapse)
+
+
+def test_random_displacement_potential_invalid_arguments_raise():
+    with pytest.raises(ValueError):
+        random_displacement_potential(4, k_min=1.0, k_max=2.0, amplitude_scale=0.1, ndim=3)
+    with pytest.raises(ValueError):
+        random_displacement_potential(0, k_min=1.0, k_max=2.0, amplitude_scale=0.1)
+    with pytest.raises(ValueError):
+        random_displacement_potential(4, k_min=2.0, k_max=1.0, amplitude_scale=0.1)
+
+
+def test_lagrangian_grid_shape_and_bounds():
+    q = lagrangian_grid(7, 3.0)
+    assert q.shape == (49, 2)
+    assert np.all(q >= 0.0)
+    assert np.all(q < 3.0)

physicskit/astro/tests/test_galactic_dynamics.py

@@ -0,0 +1,46 @@
+import numpy as np
+import pytest
+from scipy.integrate import quad
+
+from physicskit.astro.galactic_dynamics import (
+    circular_velocity,
+    nfw_density,
+    nfw_enclosed_mass,
+    nfw_potential,
+)
+
+
+def test_nfw_enclosed_mass_matches_numerical_integral_of_density():
+    rho_s, r_s = 2.0, 3.0
+    for r in [1.0, 3.0, 10.0]:
+        expected, _ = quad(lambda rp: 4 * np.pi * rp**2 * nfw_density(rp, rho_s, r_s), 0, r)
+        assert nfw_enclosed_mass(r, rho_s, r_s) == pytest.approx(expected, rel=1e-4)
+
+
+def test_nfw_enclosed_mass_monotonically_increasing():
+    rho_s, r_s = 1.0, 1.0
+    r = np.linspace(0.1, 20, 50)
+    M = nfw_enclosed_mass(r, rho_s, r_s)
+    assert np.all(np.diff(M) > 0)
+
+
+def test_nfw_potential_derivative_matches_enclosed_mass():
+    rho_s, r_s, G = 1.5, 2.0, 1.0
+    r0 = 4.0
+    dr = 1e-5
+    dPhi_dr = (nfw_potential(r0 + dr, rho_s, r_s, G) - nfw_potential(r0 - dr, rho_s, r_s, G)) / (2 * dr)
+    expected = G * nfw_enclosed_mass(r0, rho_s, r_s) / r0**2
+    assert dPhi_dr == pytest.approx(expected, rel=1e-4)
+
+
+def test_circular_velocity_point_mass_is_keplerian():
+    M0, G, r = 5.0, 1.0, 3.0
+    v = circular_velocity(r, lambda rr: M0, G=G)
+    assert v == pytest.approx(np.sqrt(G * M0 / r))
+
+
+def test_nfw_rotation_curve_is_roughly_flat_at_large_radius():
+    rho_s, r_s = 1.0, 1.0
+    r = np.linspace(5 * r_s, 30 * r_s, 50)
+    v = circular_velocity(r, lambda rr: nfw_enclosed_mass(rr, rho_s, r_s))
+    assert np.std(v) / np.mean(v) < 0.2

physicskit/astro/tests/test_nbody.py

@@ -0,0 +1,81 @@
+import numpy as np
+import pytest
+
+from physicskit.astro.nbody import NBodySystem, figure_eight_initial_conditions, gravitational_acceleration, leapfrog_step
+
+
+def test_gravitational_acceleration_two_body():
+    positions = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]])
+    masses = np.array([1.0, 1.0])
+    acc = gravitational_acceleration(positions, masses, G=1.0)
+    assert acc[0] == pytest.approx([1.0, 0.0, 0.0], abs=1e-10)
+    assert acc[1] == pytest.approx([-1.0, 0.0, 0.0], abs=1e-10)
+
+
+def test_leapfrog_step_shapes():
+    positions = np.random.default_rng(0).normal(size=(4, 3))
+    velocities = np.zeros((4, 3))
+    masses = np.ones(4)
+    new_pos, new_vel = leapfrog_step(positions, velocities, masses, dt=0.01)
+    assert new_pos.shape == (4, 3)
+    assert new_vel.shape == (4, 3)
+
+
+def _circular_two_body_system(G=1.0, m=1.0, d=2.0):
+    v = np.sqrt(G * m / (2 * d))
+    positions = np.array([[d / 2, 0.0, 0.0], [-d / 2, 0.0, 0.0]])
+    velocities = np.array([[0.0, v, 0.0], [0.0, -v, 0.0]])
+    masses = np.array([m, m])
+    return NBodySystem(positions, velocities, masses, G=G), v, d
+
+
+def test_two_body_circular_orbit_conserves_energy_and_angular_momentum():
+    system, v, d = _circular_two_body_system()
+    E0 = system.total_energy()
+    L0 = system.total_angular_momentum()
+
+    period = 2 * np.pi * (d / 2) / v
+    n_steps = 2000
+    dt = period * 5 / n_steps  # five orbits
+    system.simulate(dt, n_steps)
+
+    E1 = system.total_energy()
+    L1 = system.total_angular_momentum()
+    assert E1 == pytest.approx(E0, rel=1e-3)
+    assert L1 == pytest.approx(L0, rel=1e-3)
+
+
+def test_two_body_orbit_returns_near_starting_position_after_one_period():
+    system, v, d = _circular_two_body_system()
+    period = 2 * np.pi * (d / 2) / v
+    n_steps = 500
+    dt = period / n_steps
+    history = system.simulate(dt, n_steps)
+    assert history[-1] == pytest.approx(history[0], abs=1e-2)
+
+
+def test_figure_eight_initial_conditions_has_zero_net_momentum():
+    """The choreography's whole premise is that the three equal masses
+    share a common center of mass at rest; a nonzero net momentum would
+    mean the figure drifts rather than retracing itself in place."""
+    positions, velocities, masses = figure_eight_initial_conditions()
+    momentum = np.sum(masses[:, None] * velocities, axis=0)
+    assert momentum == pytest.approx([0.0, 0.0, 0.0], abs=1e-8)
+
+
+def test_figure_eight_initial_conditions_is_periodic_and_conserves_energy():
+    """Integrating for one full period (T ~ 6.32591398, the well-known
+    value for this choreography) should return each body close to its
+    starting position, with energy conserved by the symplectic leapfrog
+    integrator."""
+    positions, velocities, masses = figure_eight_initial_conditions()
+    system = NBodySystem(positions, velocities, masses)
+    E0 = system.total_energy()
+
+    period = 6.32591398
+    n_steps = 20000
+    dt = period / n_steps
+    system.simulate(dt, n_steps)
+
+    assert system.total_energy() == pytest.approx(E0, rel=1e-6)
+    assert system.positions == pytest.approx(positions, abs=1e-2)

physicskit/astro/tests/test_orbital_mechanics.py

@@ -0,0 +1,73 @@
+import numpy as np
+import pytest
+
+from physicskit.astro.orbital_mechanics import (
+    hohmann_transfer,
+    orbital_elements_from_state,
+    orbital_period,
+    state_from_orbital_elements,
+    vis_viva_speed,
+)
+
+
+def test_orbital_period_matches_kepler_third_law():
+    mu = 1.0
+    a = 2.0
+    T = orbital_period(a, mu)
+    assert T == pytest.approx(2 * np.pi * np.sqrt(a**3 / mu))
+
+
+def test_vis_viva_parabolic_limit():
+    v = vis_viva_speed(r=1.0, a=np.inf, mu=1.0)
+    assert v == pytest.approx(np.sqrt(2.0))
+
+
+def test_vis_viva_circular_orbit_matches_period():
+    mu = 1.0
+    r = 3.0
+    v = vis_viva_speed(r, a=r, mu=mu)
+    period_from_v = 2 * np.pi * r / v
+    assert period_from_v == pytest.approx(orbital_period(r, mu))
+
+
+ELEMENT_CASES = [
+    (1.0, 0.3, 0.5, 1.0, 0.7, 1.2),
+    (2.0, 0.1, 1.2, 2.5, 0.3, 4.0),
+    (1.5, 0.5, 0.9, 0.2, 2.1, 0.5),
+    (1.5, 0.4, 0.8, 4.0, 4.0, 1.0),  # n_vec[1]<0 and e_vec[2]<0 branches of raan/argp
+    (1.5, 0.4, 0.0, 0.0, 0.0, 1.0),  # equatorial (n<=1e-12): raan, argp forced to 0
+    (1.5, 0.0, 0.8, 1.0, 0.0, 0.0),  # circular (e<=1e-12): argp, nu forced to 0
+]
+
+
+@pytest.mark.parametrize("a,e,i,raan,argp,nu", ELEMENT_CASES)
+def test_orbital_elements_roundtrip(a, e, i, raan, argp, nu):
+    mu = 1.0
+    r_vec, v_vec = state_from_orbital_elements(a, e, i, raan, argp, nu, mu)
+    a2, e2, i2, raan2, argp2, nu2 = orbital_elements_from_state(r_vec, v_vec, mu)
+    assert a2 == pytest.approx(a, rel=1e-8)
+    assert e2 == pytest.approx(e, rel=1e-8)
+    assert i2 == pytest.approx(i, rel=1e-8)
+    assert raan2 == pytest.approx(raan % (2 * np.pi), rel=1e-8)
+    assert argp2 == pytest.approx(argp % (2 * np.pi), rel=1e-8)
+    assert nu2 == pytest.approx(nu % (2 * np.pi), rel=1e-8)
+
+
+def test_hohmann_transfer_outward_deltav_matches_speed_differences():
+    mu, r1, r2 = 1.0, 1.0, 4.0
+    a_t = (r1 + r2) / 2.0
+    dv1, dv2, t = hohmann_transfer(r1, r2, mu)
+    expected_dv1 = vis_viva_speed(r1, a_t, mu) - vis_viva_speed(r1, r1, mu)
+    expected_dv2 = vis_viva_speed(r2, r2, mu) - vis_viva_speed(r2, a_t, mu)
+    assert dv1 == pytest.approx(abs(expected_dv1))
+    assert dv2 == pytest.approx(abs(expected_dv2))
+    assert dv1 > 0
+    assert dv2 > 0
+    assert t == pytest.approx(orbital_period(a_t, mu) / 2.0)
+
+
+def test_hohmann_transfer_inward_is_positive_magnitudes():
+    dv1, dv2, t = hohmann_transfer(4.0, 1.0, mu=1.0)
+    assert dv1 > 0
+    assert dv2 > 0
+    assert t > 0