physicskit 0.1.0__py3-none-any.whl

physicskit/astro/nbody.py

@@ -0,0 +1,208 @@
+r"""Direct-summation N-body gravitational dynamics.
+
+Uses **gravitational units with** :math:`G=1` by default -- every function
+accepts ``G`` as a keyword argument, the same convention
+:mod:`physicskit.relativity` uses for its own geometrized units.
+
+- :func:`gravitational_acceleration` -- pairwise Plummer-softened gravity.
+- :func:`leapfrog_step` -- one symplectic kick-drift-kick integration step.
+- :class:`NBodySystem` -- a stateful N-body simulation with energy and
+  angular-momentum diagnostics.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+from numba import njit
+
+from physicskit.integrators.fixed_step import leapfrog_step as _leapfrog_step
+
+__all__ = ["gravitational_acceleration", "leapfrog_step", "NBodySystem", "figure_eight_initial_conditions"]
+
+
+@njit(cache=True)
+def _pairwise_acceleration(positions, masses, G, softening):
+    n = positions.shape[0]
+    acc = np.zeros((n, 3))
+    eps2 = softening**2
+    for i in range(n):
+        ax, ay, az = 0.0, 0.0, 0.0
+        for j in range(n):
+            if i == j:
+                continue
+            dx = positions[j, 0] - positions[i, 0]
+            dy = positions[j, 1] - positions[i, 1]
+            dz = positions[j, 2] - positions[i, 2]
+            dist2 = dx * dx + dy * dy + dz * dz + eps2
+            inv_dist3 = dist2 ** (-1.5)
+            f = G * masses[j] * inv_dist3
+            ax += f * dx
+            ay += f * dy
+            az += f * dz
+        acc[i, 0] = ax
+        acc[i, 1] = ay
+        acc[i, 2] = az
+    return acc
+
+
+def gravitational_acceleration(positions, masses, G=1.0, softening=0.0):
+    r"""Pairwise Plummer-softened gravitational acceleration.
+
+    .. math::
+
+        \vec a_i = G\sum_{j\neq i} m_j\,
+        \frac{\vec r_j-\vec r_i}{\left(|\vec r_j-\vec r_i|^2+\epsilon^2\right)^{3/2}}
+
+    Parameters
+    ----------
+    positions : ndarray of shape (N, 3)
+    masses : ndarray of shape (N,)
+    G : float, default=1.0
+    softening : float, default=0.0
+        Plummer softening length :math:`\epsilon`.
+
+    Returns
+    -------
+    ndarray of shape (N, 3)
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> pos = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]])
+    >>> m = np.array([1.0, 1.0])
+    >>> acc = gravitational_acceleration(pos, m)
+    >>> round(float(acc[0, 0]), 6)
+    1.0
+    """
+    positions = np.ascontiguousarray(positions, dtype=np.float64)
+    masses = np.ascontiguousarray(masses, dtype=np.float64)
+    return _pairwise_acceleration(positions, masses, float(G), float(softening))
+
+
+@njit(cache=True)
+def _nbody_accel(positions, t, params):
+    """Adapt :func:`_pairwise_acceleration` to the shared ``force(pos, t,
+    params) -> acc`` convention, packing ``masses``, ``G`` and ``softening``
+    into a single ``params`` vector (``masses`` followed by ``[G,
+    softening]``) so it can be passed as a first-class function into
+    :func:`physicskit.integrators.leapfrog_step`."""
+    n = positions.shape[0]
+    masses = params[:n]
+    G = params[n]
+    softening = params[n + 1]
+    return _pairwise_acceleration(positions, masses, G, softening)
+
+
+def leapfrog_step(positions, velocities, masses, dt, G=1.0, softening=0.0):
+    """One kick-drift-kick leapfrog (symplectic) integration step.
+
+    Delegates to the shared :func:`physicskit.integrators.leapfrog_step`.
+
+    Parameters
+    ----------
+    positions : ndarray of shape (N, 3)
+    velocities : ndarray of shape (N, 3)
+    masses : ndarray of shape (N,)
+    dt : float
+        Timestep.
+    G : float, default=1.0
+    softening : float, default=0.0
+
+    Returns
+    -------
+    new_positions, new_velocities : ndarray of shape (N, 3)
+    """
+    positions = np.ascontiguousarray(positions, dtype=np.float64)
+    velocities = np.ascontiguousarray(velocities, dtype=np.float64)
+    masses = np.ascontiguousarray(masses, dtype=np.float64)
+    params = np.concatenate((masses, np.array([G, softening], dtype=np.float64)))
+    return _leapfrog_step(_nbody_accel, positions, velocities, 0.0, dt, params)
+
+
+class NBodySystem:
+    """A stateful direct-summation N-body gravitational simulation.
+
+    Parameters
+    ----------
+    positions : ndarray of shape (N, 3)
+    velocities : ndarray of shape (N, 3)
+    masses : ndarray of shape (N,)
+    G : float, default=1.0
+    softening : float, default=0.0
+    """
+
+    def __init__(self, positions, velocities, masses, G=1.0, softening=0.0):
+        self.positions = np.array(positions, dtype=np.float64, copy=True)
+        self.velocities = np.array(velocities, dtype=np.float64, copy=True)
+        self.masses = np.array(masses, dtype=np.float64, copy=True)
+        self.G = G
+        self.softening = softening
+
+    def step(self, dt):
+        """Advance the system's stored state by one leapfrog step, in place."""
+        self.positions, self.velocities = leapfrog_step(self.positions, self.velocities, self.masses, dt, self.G, self.softening)
+
+    def simulate(self, dt, n_steps):
+        """Advance ``n_steps`` steps, returning the position history.
+
+        Returns
+        -------
+        ndarray of shape (n_steps + 1, N, 3)
+        """
+        n = self.positions.shape[0]
+        history = np.zeros((n_steps + 1, n, 3))
+        history[0] = self.positions
+        for k in range(n_steps):
+            self.step(dt)
+            history[k + 1] = self.positions
+        return history
+
+    def total_energy(self):
+        r"""Kinetic plus (softened) gravitational potential energy, at the current state."""
+        kinetic = 0.5 * np.sum(self.masses[:, None] * self.velocities**2)
+        n = self.positions.shape[0]
+        potential = 0.0
+        eps2 = self.softening**2
+        for i in range(n):
+            for j in range(i + 1, n):
+                d = np.linalg.norm(self.positions[i] - self.positions[j])
+                potential -= self.G * self.masses[i] * self.masses[j] / np.sqrt(d**2 + eps2)
+        return float(kinetic + potential)
+
+    def total_angular_momentum(self):
+        r"""Total angular momentum, :math:`\vec L=\sum_im_i\,\vec r_i\times\vec v_i`."""
+        L = np.sum(self.masses[:, None] * np.cross(self.positions, self.velocities), axis=0)
+        return L
+
+
+def figure_eight_initial_conditions():
+    r"""Initial conditions for the equal-mass planar figure-eight three-body choreography.
+
+    Three equal masses, in :math:`G=1` units, chase each other forever
+    around a single figure-eight-shaped curve -- a periodic solution
+    discovered numerically by Moore (1993) and proven to exist by
+    Chenciner & Montgomery (2000, *Annals of Mathematics* 152). Unlike
+    typical three-body configurations it is planar, collision-free, and
+    exactly periodic, making it a clean, visually striking demo for
+    :class:`NBodySystem`.
+
+    Returns
+    -------
+    positions : ndarray of shape (3, 3)
+    velocities : ndarray of shape (3, 3)
+    masses : ndarray of shape (3,)
+
+    Examples
+    --------
+    >>> positions, velocities, masses = figure_eight_initial_conditions()
+    >>> system = NBodySystem(positions, velocities, masses)
+    >>> history = system.simulate(dt=0.001, n_steps=100)
+    >>> history.shape
+    (101, 3, 3)
+    """
+    r1 = np.array([0.9700436, -0.24308753, 0.0])
+    v3 = np.array([0.93240737, 0.86473146, 0.0])
+    positions = np.array([r1, -r1, [0.0, 0.0, 0.0]])
+    velocities = np.array([-0.5 * v3, -0.5 * v3, v3])
+    masses = np.array([1.0, 1.0, 1.0])
+    return positions, velocities, masses

physicskit/astro/orbital_mechanics.py

@@ -0,0 +1,220 @@
+r"""Two-body Keplerian orbital mechanics: elements, state vectors, and transfers.
+
+Uses **gravitational units with** :math:`G=1` by default; the standard
+gravitational parameter :math:`\mu=GM` appears directly as a function
+argument, the normal convention in orbital mechanics regardless of unit
+system (the same style :mod:`physicskit.relativity` uses for its own
+geometrized units).
+
+- :func:`orbital_period`, :func:`vis_viva_speed` -- Kepler's third law and
+  the vis-viva equation.
+- :func:`orbital_elements_from_state`, :func:`state_from_orbital_elements`
+  -- converting between Cartesian state vectors and the six classical
+  orbital elements.
+- :func:`hohmann_transfer` -- the two-burn transfer between circular orbits.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+__all__ = [
+    "orbital_period",
+    "vis_viva_speed",
+    "orbital_elements_from_state",
+    "state_from_orbital_elements",
+    "hohmann_transfer",
+]
+
+
+def orbital_period(a, mu):
+    r"""Kepler's third law, :math:`T=2\pi\sqrt{a^3/\mu}`.
+
+    Parameters
+    ----------
+    a : float
+        Semi-major axis (bound orbit, ``a > 0``).
+    mu : float
+        Gravitational parameter, :math:`\mu=GM`.
+
+    Returns
+    -------
+    float
+
+    Examples
+    --------
+    >>> round(orbital_period(1.0, 1.0) / (2 * np.pi), 6)
+    1.0
+    """
+    return float(2.0 * np.pi * np.sqrt(a**3 / mu))
+
+
+def vis_viva_speed(r, a, mu):
+    r"""The vis-viva equation, :math:`v=\sqrt{\mu\left(2/r-1/a\right)}`.
+
+    Works for ``a > 0`` (ellipse), ``a < 0`` (hyperbola), and ``a = inf``
+    (parabolic escape trajectory, where the formula reduces to
+    :math:`v=\sqrt{2\mu/r}`).
+
+    Parameters
+    ----------
+    r : float
+        Current orbital radius.
+    a : float
+        Semi-major axis.
+    mu : float
+        Gravitational parameter.
+
+    Returns
+    -------
+    float
+
+    Examples
+    --------
+    >>> round(vis_viva_speed(1.0, np.inf, 1.0), 6)
+    1.414214
+    """
+    return float(np.sqrt(mu * (2.0 / r - 1.0 / a)))
+
+
+def orbital_elements_from_state(r_vec, v_vec, mu):
+    r"""Classical (Keplerian) orbital elements from a Cartesian state vector.
+
+    Returns semi-major axis, eccentricity, inclination, right ascension of
+    the ascending node, argument of periapsis, and true anomaly (angles
+    in radians). For a near-equatorial orbit (:math:`i\approx0` or
+    :math:`\pi`), the node vector is near zero and ``raan`` is undefined;
+    this implementation returns ``raan=0.0`` in that case rather than
+    raising -- a documented simplification, not a fully general
+    equatorial-orbit treatment.
+
+    Parameters
+    ----------
+    r_vec, v_vec : ndarray of shape (3,)
+        Position and velocity.
+    mu : float
+        Gravitational parameter.
+
+    Returns
+    -------
+    a, e, i, raan, argp, nu : float
+
+    Examples
+    --------
+    >>> r_vec, v_vec = state_from_orbital_elements(1.0, 0.3, 0.5, 1.0, 0.7, 1.2, mu=1.0)
+    >>> a, e, i, raan, argp, nu = orbital_elements_from_state(r_vec, v_vec, mu=1.0)
+    >>> [round(x, 6) for x in (a, e, i, raan, argp, nu)]
+    [1.0, 0.3, 0.5, 1.0, 0.7, 1.2]
+    """
+    r_vec = np.asarray(r_vec, dtype=float)
+    v_vec = np.asarray(v_vec, dtype=float)
+    r = np.linalg.norm(r_vec)
+    v = np.linalg.norm(v_vec)
+
+    h_vec = np.cross(r_vec, v_vec)
+    h = np.linalg.norm(h_vec)
+    n_vec = np.cross([0.0, 0.0, 1.0], h_vec)
+    n = np.linalg.norm(n_vec)
+    e_vec = np.cross(v_vec, h_vec) / mu - r_vec / r
+    e = np.linalg.norm(e_vec)
+
+    energy = v**2 / 2.0 - mu / r
+    a = np.inf if abs(e - 1.0) < 1e-12 else -mu / (2.0 * energy)
+
+    i = float(np.arccos(np.clip(h_vec[2] / h, -1.0, 1.0)))
+
+    if n > 1e-12:
+        raan = float(np.arccos(np.clip(n_vec[0] / n, -1.0, 1.0)))
+        if n_vec[1] < 0:
+            raan = 2 * np.pi - raan
+    else:
+        raan = 0.0
+
+    if n > 1e-12 and e > 1e-12:
+        argp = float(np.arccos(np.clip(np.dot(n_vec, e_vec) / (n * e), -1.0, 1.0)))
+        if e_vec[2] < 0:
+            argp = 2 * np.pi - argp
+    else:
+        argp = 0.0
+
+    if e > 1e-12:
+        nu = float(np.arccos(np.clip(np.dot(e_vec, r_vec) / (e * r), -1.0, 1.0)))
+        if np.dot(r_vec, v_vec) < 0:
+            nu = 2 * np.pi - nu
+    else:
+        nu = 0.0
+
+    return float(a), float(e), i, raan, argp, nu
+
+
+def state_from_orbital_elements(a, e, i, raan, argp, nu, mu):
+    """The inverse of :func:`orbital_elements_from_state`.
+
+    Parameters
+    ----------
+    a, e, i, raan, argp, nu : float
+        Classical orbital elements (angles in radians).
+    mu : float
+        Gravitational parameter.
+
+    Returns
+    -------
+    r_vec, v_vec : ndarray of shape (3,)
+    """
+    p = a * (1.0 - e**2)
+    r = p / (1.0 + e * np.cos(nu))
+    r_pf = np.array([r * np.cos(nu), r * np.sin(nu), 0.0])
+    v_pf = np.sqrt(mu / p) * np.array([-np.sin(nu), e + np.cos(nu), 0.0])
+
+    cO, sO = np.cos(raan), np.sin(raan)
+    ci, si = np.cos(i), np.sin(i)
+    cw, sw = np.cos(argp), np.sin(argp)
+
+    # 3-1-3 Euler rotation (raan, i, argp): perifocal -> reference frame.
+    R = np.array(
+        [
+            [cO * cw - sO * sw * ci, -cO * sw - sO * cw * ci, sO * si],
+            [sO * cw + cO * sw * ci, -sO * sw + cO * cw * ci, -cO * si],
+            [sw * si, cw * si, ci],
+        ]
+    )
+    return R @ r_pf, R @ v_pf
+
+
+def hohmann_transfer(r1, r2, mu):
+    r"""The two burns and transfer time for a Hohmann transfer.
+
+    Parameters
+    ----------
+    r1 : float
+        Initial circular orbit radius.
+    r2 : float
+        Final circular orbit radius.
+    mu : float
+        Gravitational parameter.
+
+    Returns
+    -------
+    delta_v1 : float
+        Speed change (magnitude) to leave the initial circular orbit and
+        enter the transfer ellipse.
+    delta_v2 : float
+        Speed change (magnitude) to circularize at ``r2``.
+    transfer_time : float
+        Time to complete the transfer, half the transfer ellipse's period.
+
+    Examples
+    --------
+    >>> dv1, dv2, t = hohmann_transfer(1.0, 4.0, mu=1.0)
+    >>> dv1 > 0 and dv2 > 0 and t > 0
+    True
+    """
+    a_t = (r1 + r2) / 2.0
+    v1_circular = vis_viva_speed(r1, r1, mu)
+    v1_transfer = vis_viva_speed(r1, a_t, mu)
+    v2_transfer = vis_viva_speed(r2, a_t, mu)
+    v2_circular = vis_viva_speed(r2, r2, mu)
+    delta_v1 = abs(v1_transfer - v1_circular)
+    delta_v2 = abs(v2_circular - v2_transfer)
+    transfer_time = orbital_period(a_t, mu) / 2.0
+    return float(delta_v1), float(delta_v2), float(transfer_time)