orbit-gen 0.1.0__py3-none-any.whl

orbit_gen/__init__.py

@@ -0,0 +1,3 @@
+from orbit_gen.orbit import Orbit
+from orbit_gen.family import Family
+from orbit_gen.setup.config import RunConfig

orbit_gen/family.py

@@ -0,0 +1,296 @@
+import numpy as np
+import spiceypy as spice
+from matplotlib import pyplot as plt
+from scipy.optimize import fsolve
+import astropy.units as u
+
+from orbit_gen.utils import eom, units
+from orbit_gen.orbit import Orbit
+from orbit_gen.setup.config import RunConfig
+from orbit_gen.setup.system_config import CanonicalUnits
+
+
+class Family:
+    """Family of CR3BP periodic orbits
+
+    Generates a family of periodic orbits in the circular restricted
+    three-body problem (CR3BP).
+
+    Args:
+        cfg (RunConfig):
+            Full run configuration.
+
+    Attributes:
+        mu_star (float):
+            CR3BP mass ratio.
+        m1 (float):
+            Normalized primary mass (1 - mu_star).
+        m2 (float):
+            Normalized secondary mass (mu_star).
+        r_secondary (float):
+            Radius of the secondary body in canonical distance units (DU).
+        orbits (list[Orbit]):
+            Collection of valid periodic orbit solutions obtained.
+    """
+
+    def __init__(self, cfg: RunConfig):
+        self.cfg = cfg
+        self.orbits: list[Orbit] = []
+
+        self._setup_constants()
+        self._run_continuation()
+
+    def _setup_constants(self):
+        """Load SPICE kernel and compute CR3BP constants from system config."""
+        for path in self.cfg.kernel_paths.values():
+            spice.furnsh(str(path))
+
+        sys = self.cfg.system
+        gm_primary = spice.bodvrd(sys.primary, "GM", 1)[1][0]
+        gm_secondary = spice.bodvrd(sys.secondary, "GM", 1)[1][0]
+
+        self.mu_star = gm_secondary / (gm_primary + gm_secondary)
+        self.m1 = 1 - self.mu_star
+        self.m2 = self.mu_star
+
+        dist_km = sys.canonical_units.dist_km
+        time_seconds = np.sqrt(dist_km**3 / (gm_primary + gm_secondary))
+        sys.canonical_units = CanonicalUnits(
+            dist_km=dist_km, time_days=time_seconds / 86400
+        )
+
+        radii = spice.bodvrd(sys.secondary, "RADII", 3)[1][0]
+        self.r_secondary = units.pos_to_canonical(radii * u.km, sys.canonical_units)
+
+    def _run_continuation(self):
+        """Outer pseudo-arclength loop. Populates self.orbits."""
+        cont = self.cfg.continuation
+        Tp_lim = units.time_to_canonical(
+            cont.max_period_days * u.d, self.cfg.system.canonical_units
+        )
+        stm_identity = np.reshape(np.eye(6), (1, 36))[0]
+
+        # Free variables: [x0, z0, vy0, T/2]
+        X = self.cfg.ic.as_list
+        z = np.array([0.0, 0.0, 0.0, -1.0])
+
+        while X[-1] * 2 < Tp_lim and len(self.orbits) < cont.max_solutions:
+            X, error = self._correct_orbit(X, stm_identity)
+
+            if X[-1] < 0:
+                print("Negative period.")
+                break
+            if error > cont.eps:
+                print("Corrector did not converge.")
+                break
+
+            period_days = units.time_to_dimensional(
+                2 * X[-1], self.cfg.system.canonical_units
+            ).to_value(u.day)
+            print(f"Period: {period_days:.4f} days")
+
+            states, times, valid = self._propagate_and_validate(X)
+
+            if valid:
+                orbit = Orbit(
+                    x0=X[0],
+                    z0=X[1],
+                    vy0=X[2],
+                    period=times[-1],
+                    mu_star=self.mu_star,
+                    canonical_units=self.cfg.system.canonical_units,
+                )
+                # Store already-computed trajectory so propagate() isn't needed again
+                orbit.states = states
+                orbit.times = times
+                self.orbits.append(orbit)
+
+            X, z, success = self._continuation_step(X, z)
+            if not success:
+                break
+
+            print(f"Orbits found: {len(self.orbits)}\n")
+
+        if not self.orbits:
+            raise RuntimeError("Continuation failed")
+
+    def _correct_orbit(self, X: list, stm_identity: np.ndarray) -> tuple[list, float]:
+        """Single-shooting differential corrector
+
+        Newton-style differential correction using a single-shooting formulation
+        for periodic orbit boundary condition enforcement.
+
+        Args:
+            X (array-like):
+                Free-variable vector of the form [x0, z0, vy0, T/2].
+            stm_identity (ndarray):
+                Flattened 6x6 identity matrix used to initialize/propagate the
+                state transition matrix.
+
+        Returns:
+            X (ndarray):
+                Corrected free-variable vector after convergence.
+            error (float):
+                Final residual norm ||F(X)|| at termination.
+        """
+
+        cont = self.cfg.continuation
+        error = np.inf
+
+        for i in range(cont.max_iter):
+            X_full = np.append(X, stm_identity).tolist()
+
+            Fx, state_half_T, Phi = eom.constraint(X_full, self.mu_star)
+            error = np.linalg.norm(Fx)
+
+            if error < cont.eps:
+                break
+
+            dFx = eom.jacobian(X, self.mu_star, state_half_T, Phi)
+            X = X - 0.3 * dFx.T @ (np.linalg.inv(dFx @ dFx.T) @ Fx)
+
+            print(
+                f"  iter {i+1:2d} | error = {error:.3e} | cond = {np.linalg.cond(dFx @ dFx.T):.2e}"
+            )
+
+        return X, error
+
+    def _propagate_and_validate(self, X: list) -> tuple[np.ndarray, np.ndarray, bool]:
+        """Propagate one full orbital period and check secondary-body intersection
+
+        Numerically integrates the trajectory over a full orbital period from the
+        given initial conditions, then evaluates whether the trajectory intersects
+        or enters the secondary body's exclusion region.
+
+        Args:
+            X (array-like):
+                Free-variable vector [x0, z0, vy0, T/2] defining the initial state
+                and half-period used for symmetric propagation.
+
+        Returns:
+            states (ndarray):
+                Propagated state history with shape (N, 6), where each row is
+                [x, y, z, vx, vy, vz].
+            times (ndarray):
+                Time samples corresponding to each state, shape (N,).
+            valid (bool):
+                False if the trajectory intersects the secondary body or violates
+                the exclusion radius; True otherwise.
+        """
+
+        X_full = X.copy()
+        X_full[-1] = 2 * X_full[-1]
+
+        states, times = eom.propagate(X_full, self.mu_star)
+        positions = states[:, 0:3]
+        r_mag = np.linalg.norm(positions, axis=1)
+
+        periapsis_km = units.pos_to_dimensional(
+            min(r_mag), self.cfg.system.canonical_units
+        ).to_value(u.km)
+        print(f"  Periapsis: {periapsis_km:.2f} km")
+
+        valid = not np.any(r_mag < self.r_secondary)
+        if not valid:
+            print("  Intersects secondary — discarding.")
+
+        return states, times, valid
+
+    def _continuation_step(
+        self, X: list, z: np.ndarray
+    ) -> tuple[list, np.ndarray, bool]:
+        """Advance one pseudo-arclength continuation step along an orbit family
+
+        Performs a predictor-corrector step along the solution manifold using a
+        pseudo-arclength continuation method. Advances the current solution estimate
+        and updates the tangent direction along the family.
+
+        Args:
+            X (ndarray):
+                Current free-variable vector at the current solution point.
+            z (ndarray):
+                Current tangent direction vector along the continuation manifold.
+
+        Returns:
+            X_new (ndarray):
+                Predicted/corrected free-variable vector for the next orbit.
+            z_new (ndarray):
+                Updated tangent direction at the new solution point.
+            success (bool):
+                False if the Jacobian is singular or the correction step fails;
+                True otherwise.
+        """
+
+        sol_predictor = X + z * self.cfg.continuation.step
+
+        fss = fsolve(
+            eom.continuation_constraint,
+            X,
+            args=(z, sol_predictor, self.mu_star),
+            full_output=True,
+            xtol=1e-6,
+        )
+        X_new = fss[0]
+
+        # Reconstruct Jacobian
+        Q = fss[1]["fjac"]
+        Rs = fss[1]["r"]
+        R = np.zeros((4, 4))
+        idx, col = np.triu_indices(4, k=0)
+        R[idx, col] = Rs
+        J = Q.T @ R
+
+        try:
+            z_new = np.linalg.inv(J) @ z
+            z_new = z_new / np.linalg.norm(z_new)
+            return X_new, z_new, True
+        except np.linalg.LinAlgError:
+            print("Singular Jacobian — stopping continuation.")
+            return X_new, z, False
+
+    def plot(self):
+        """Plot all orbits in the family on a single 3D figure
+
+        Generates a 3D visualization of all stored periodic orbits in the family,
+        typically in canonical CR3BP coordinates. Each orbit is plotted in the same
+        figure for qualitative comparison of geometry across the continuation family.
+
+        Args:
+            show (bool):
+                If True, displays the figure via matplotlib.pyplot.show().
+                If False, the figure is created but not rendered interactively.
+        """
+
+        fig = plt.figure()
+        ax = fig.add_subplot(projection="3d")
+
+        for orbit in self.orbits:
+            orbit.plot(ax=ax)
+
+        ax.set_title(f"{self.cfg.ic.name} family — {len(self.orbits)} orbits")
+
+        plt.show()
+
+        return ax
+
+    def save(self):
+        """Save all orbit ICs, periods, and mu_star to cfg.output_path."""
+        self.cfg.output_dir.mkdir(parents=True, exist_ok=True)
+
+        states = np.array([o.states[0, :] for o in self.orbits])
+        periods = np.array([o.period for o in self.orbits])
+
+        np.savez(
+            self.cfg.output_path,
+            states=states,
+            periods=periods,
+            mu_star=self.mu_star,
+        )
+        print(f"Saved {len(self.orbits)} orbits to {self.cfg.output_path}")
+
+    def __repr__(self):
+        return (
+            f"Family(system={self.cfg.system.name}, "
+            f"ic={self.cfg.ic.name}, "
+            f"orbits={len(self.orbits)})"
+        )

orbit_gen/orbit.py

@@ -0,0 +1,135 @@
+import numpy as np
+from matplotlib import pyplot as plt
+
+from orbit_gen.utils import eom, units
+from orbit_gen.setup.system_config import CanonicalUnits
+import astropy.units as u
+
+
+class Orbit:
+    """Single corrected periodic orbit in the CR3BP
+
+    Represents a converged periodic orbit obtained from a differential correction
+    procedure in the circular restricted three-body problem (CR3BP). The orbit
+    is defined at a symmetry-plane crossing and can be numerically propagated to
+    recover the full trajectory.
+
+    Args:
+        x0 (float):
+            x-position at symmetry-plane crossing in canonical distance units (DU).
+        z0 (float):
+            z-position at symmetry-plane crossing in DU.
+        vy0 (float):
+            y-velocity at symmetry-plane crossing in canonical units (DU/TU).
+        period (float):
+            Full orbital period T in canonical time units (TU).
+        mu_star (float):
+            CR3BP mass ratio (normalized secondary mass parameter).
+
+    Attributes:
+        states (ndarray or None):
+            Trajectory array of shape (N, 6) containing
+            [x, y, z, vx, vy, vz] along the orbit. None until `propagate()` is called.
+        times (ndarray or None):
+            Time array of shape (N,) corresponding to `states`. None until
+            `propagate()` is called.
+    """
+
+    def __init__(
+        self,
+        x0: float,
+        z0: float,
+        vy0: float,
+        period: float,
+        mu_star: float,
+        canonical_units: CanonicalUnits,
+    ):
+        self.x0 = x0
+        self.z0 = z0
+        self.vy0 = vy0
+        self.period = period
+        self.mu_star = mu_star
+        self.canonical_units = canonical_units
+
+        self.states = None
+        self.times = None
+
+    @property
+    def period_days(self) -> float:
+        """Full orbital period in days."""
+        return units.time_to_dimensional(self.period, self.canonical_units).to_value(
+            u.day
+        )
+
+    @property
+    def _free_var(self) -> list:
+        """Full-period free variable vector [x0, z0, vy0, T]."""
+        return [self.x0, self.z0, self.vy0, self.period]
+
+    def propagate(self):
+        """Integrate the trajectory for one full period.
+
+        Stores results in self.states and self.times.
+        """
+        self.states, self.times = eom.propagate(self._free_var, self.mu_star)
+
+    def plot(self, ax=None, **kwargs):
+        """Plot the orbit in 3D
+
+        Plots the trajectory of the orbit in a 3D matplotlib figure using the
+        stored state history.
+
+        Args:
+            ax (matplotlib.axes._subplots.Axes3DSubplot, optional):
+                Existing 3D axes object. If None, a new figure and 3D axes are
+                created internally.
+            **kwargs:
+                Additional keyword arguments forwarded directly to `ax.plot()`,
+                controlling styling (e.g., color, linewidth, alpha).
+
+        Returns:
+            ax (matplotlib.axes._subplots.Axes3DSubplot):
+                The matplotlib 3D axes containing the plotted orbit.
+        """
+
+        if self.states is None:
+            self.propagate()
+
+        if ax is None:
+            fig = plt.figure()
+            ax = fig.add_subplot(projection="3d")
+
+        ax.plot(self.states[:, 0], self.states[:, 1], self.states[:, 2], **kwargs)
+        ax.set_xlabel("X [DU]")
+        ax.set_ylabel("Y [DU]")
+        ax.set_zlabel("Z [DU]")
+
+        return ax
+
+    def save(self, path):
+        """Save orbit initial conditions and period to an NPZ file
+
+        Serializes the orbit’s defining parameters (initial state and period)
+        to a NumPy `.npz` archive for later reuse or reconstruction.
+
+        Args:
+            path (str or pathlib.Path):
+                Destination file path for the `.npz` archive. If the extension is
+                omitted, `.npz` is assumed.
+        """
+
+        if self.states is None:
+            self.propagate()
+
+        np.savez(
+            path,
+            ic=self.states[0, :],
+            period=self.period,
+            mu_star=self.mu_star,
+        )
+
+    def __repr__(self):
+        return (
+            f"Orbit(x0={self.x0:.6f}, z0={self.z0:.6f}, "
+            f"vy0={self.vy0:.6f}, T={self.period:.6f})"
+        )

orbit_gen/setup/config.py

@@ -0,0 +1,114 @@
+from dataclasses import dataclass, field
+from .orbit_ic import OrbitIC
+from .system_config import SystemConfig, EM_SYSTEM, SE_SYSTEM
+from .presets import PRESETS
+from .kernels import ensure_kernels
+from pathlib import Path
+import yaml
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+
+_SYSTEMS = {
+    "EM": EM_SYSTEM,
+    "SE": SE_SYSTEM,
+}
+
+
+@dataclass
+class ContinuationConfig:
+    """Parameters controlling differential correction and family continuation.
+
+    Attributes:
+        eps (float):
+            Convergence tolerance on the constraint residual norm
+        max_iter (int):
+            Maximum number of Newton iterations for a single
+            differential correction attempt
+        max_period_days (float):
+            Maximum allowable full orbit period, in days
+        max_solutions (int):
+            Maximum number of converged periodic orbits to generate before
+            terminating continuation
+        step (float):
+            Pseudo-arclength continuation step size in canonical units
+
+    """
+
+    max_iter: int = 1000
+    eps: float = 1e-6
+    step: float = 0.01
+    max_solutions: int = 10
+    max_period_days: float = 1000.0
+
+
+@dataclass
+class RunConfig:
+    """Top-level configuration for a continuation run.
+
+    Attributes:
+        continuation (ContinuationConfig):
+            Differential correction and continuation parameters
+        ic (OrbitIC):
+            Periodic orbit used to initialize continuation
+        output_dir (Path):
+            Directory where output files will be written
+        output_filename (str):
+            Name of the output .npz file containing continuation results
+        show_plots (bool):
+            If True, display plots interactively during execution
+        spice_kernel (Path):
+            Path to the SPICE meta-kernel furnished before propagation and
+            analysis
+        system (SystemConfig):
+            CR3BP system used
+
+    """
+
+    system: SystemConfig
+    ic: OrbitIC
+    continuation: ContinuationConfig = field(default_factory=ContinuationConfig)
+    output_dir: Path = PROJECT_ROOT.parent.parent / "results"
+    show_plots: bool = True
+    kernel_paths: dict[str, Path] = field(
+        init=False, default_factory=ensure_kernels, repr=False
+    )
+
+    @property
+    def output_filename(self) -> str:
+        return f"{self.ic.name}.npz"
+
+    @property
+    def output_path(self) -> Path:
+        return self.output_dir / self.output_filename
+
+    @classmethod
+    def from_yaml(cls, yaml_path: Path) -> "RunConfig":
+        with open(yaml_path) as yp:
+            data = yaml.safe_load(yp)
+
+        system = _SYSTEMS[data["system"]]
+        ic = PRESETS[data["preset"]]
+
+        if system.name == "EM":
+            period_max_default = 50.0
+        elif system.name == "SE":
+            period_max_default = 500.0
+        else:
+            raise RuntimeError("Unknown system name")
+
+        continuation = ContinuationConfig(
+            max_iter=data.get("max_iter", 100),
+            eps=data.get("eps", 1e-6),
+            step=data.get("step", 0.01),
+            max_solutions=data.get("max_solutions", 20),
+            max_period_days=data.get("max_period_days", period_max_default),
+        )
+
+        return cls(
+            system=system,
+            ic=ic,
+            continuation=continuation,
+            output_dir=PROJECT_ROOT.parent.parent
+            / Path(data.get("output_dir", "results")),
+            show_plots=data.get("show_plots", True),
+        )

orbit_gen/setup/kernels.py

@@ -0,0 +1,56 @@
+from dataclasses import dataclass
+from pathlib import Path
+from urllib.request import urlopen
+from urllib.error import URLError
+from platformdirs import user_cache_dir
+
+
+@dataclass(frozen=True)
+class KernelSpec:
+    filename: str
+    url: str
+
+
+REQUIRED_KERNELS: tuple[KernelSpec, ...] = (
+    KernelSpec(
+        filename="naif0012.tls",
+        url="https://naif.jpl.nasa.gov/pub/naif/generic_kernels/lsk/naif0012.tls",
+    ),
+    KernelSpec(
+        filename="gm_de440.tpc",
+        url="https://naif.jpl.nasa.gov/pub/naif/generic_kernels/pck/gm_de440.tpc",
+    ),
+    KernelSpec(
+        filename="pck00011.tpc",
+        url="https://naif.jpl.nasa.gov/pub/naif/generic_kernels/pck/pck00011.tpc",
+    ),
+)
+
+
+def cache_dir() -> Path:
+    direc = Path(user_cache_dir("orbit-gen")) / "kernels"
+    direc.mkdir(parents=True, exist_ok=True)
+    return direc
+
+
+def ensure_kernel(spec: KernelSpec) -> Path:
+    cache = cache_dir()
+    dest = cache / spec.filename
+
+    if dest.exists() and dest.stat().st_size > 0:
+        return dest
+
+    tmp = dest.with_suffix(dest.suffix + ".part")
+    try:
+        with urlopen(spec.url, timeout=30) as resp, open(tmp, "wb") as f:
+            f.write(resp.read())
+    except URLError as e:
+        raise RuntimeError(
+            f"Failed to download {spec.filename} from {spec.url}: {e}"
+        ) from e
+    tmp.rename(dest)
+    return dest
+
+
+def ensure_kernels() -> dict[str, Path]:
+    return {spec.filename: ensure_kernel(spec) for spec in REQUIRED_KERNELS}

orbit_gen/setup/orbit_ic.py

@@ -0,0 +1,41 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class OrbitIC:
+    """A single periodic orbit initial condition.
+
+    Attributes:
+        half_period (float):
+            Half-period T/2 in canonical time units (TU)
+        name (str):
+            Label for the orbit family, including the system and orbit type
+        vy0 (float):
+            y-velocity in canonical velocity units (DU/TU)
+        x0 (float):
+            x-position in canonical distance units (DU)
+        z0 (float):
+            z-position in canonical distance units (DU)
+
+    .. note::
+
+        y0, vx0, and vz0 are always zero by the CR3BP symmetry condition. System
+        is denoted by "EM" for Earth/Moon, and "SE" for Sun/Earth.
+
+    """
+
+    name: str
+    x0: float
+    z0: float
+    vy0: float
+    half_period: float
+
+    @property
+    def as_list(self) -> list:
+        """Return free variables as [x0, z0, vy0, T/2]."""
+        return [self.x0, self.z0, self.vy0, self.half_period]
+
+    @property
+    def full_ic(self) -> list:
+        """Return the full 7-element IC [x0, 0, z0, 0, vy0, 0, T/2]."""
+        return [self.x0, 0.0, self.z0, 0.0, self.vy0, 0.0, self.half_period]

orbit_gen/setup/presets.py

@@ -0,0 +1,102 @@
+from .orbit_ic import OrbitIC
+
+PRESETS: dict[str, OrbitIC] = {
+    "EM_L2_Northern_Halo": OrbitIC(
+        name="EM_L2_Northern",
+        x0=1.0110350588,
+        z0=-0.1731500000,
+        vy0=-0.0780141199,
+        half_period=1.3632096570 / 2,
+    ),
+    "EM_DRO_1": OrbitIC(
+        name="EM_DRO_1",
+        x0=0.583856747,
+        z0=0.0,
+        vy0=0.96455414,
+        half_period=5.70245716 / 2,
+    ),
+    "EM_DRO_2": OrbitIC(
+        name="EM_DRO_2",
+        x0=0.429519110229904,
+        z0=0.0,
+        vy0=1.440796689672539,
+        half_period=3.051133070334277,
+    ),
+    "EM_L1": OrbitIC(
+        name="EM_L1",
+        x0=0.856382122325864,
+        z0=0.181519309916197,
+        vy0=0.257898218422393,
+        half_period=1.22727308466325,
+    ),
+    "EM_Butterfly": OrbitIC(
+        name="EM_Butterfly",
+        x0=1.06896234204296,
+        z0=0.159599443574046,
+        vy0=-0.00769167653854165,
+        half_period=1.66142030228280,
+    ),
+    "EM_Unnamed_1": OrbitIC(
+        name="EM_Unnamed_1",
+        x0=0.95571113,
+        z0=0.16892834,
+        vy0=0.29101955,
+        half_period=6.8828406 / 2,
+    ),
+    "SE_L2_Northern_Halo": OrbitIC(
+        name="SE_L2_Northern_Halo",
+        x0=1.0082893998277027,
+        z0=0.011434182696839030,
+        vy0=-0.015179690420484332,
+        half_period=2.7047531095113841 / 2,
+    ),
+    "SE_L1_Vertical": OrbitIC(
+        name="SE_L1_Vertical",
+        x0=1.0022515716338869,
+        z0=-1.9346947758385478e-13,
+        vy0=-1.3010231214778553e-1,
+        half_period=6.2826376084648130 / 2,
+    ),
+    "SE_L2_Lyapunov": OrbitIC(
+        name="SE_L2_Lyapunov",
+        x0=1.0030516289641727,
+        z0=1.7613075082610678e-17,
+        vy0=4.2071729820709863e-2,
+        half_period=4.1379689311747248 / 2,
+    ),
+    "SE_L2_Small_Lyapunov": OrbitIC(
+        name="SE_L2_Small_Lyapunov",
+        x0=1.0173058431158761,
+        z0=-8.9129284257995943e-27,
+        vy0=-3.5071329991448227e-2,
+        half_period=5.5314456810353905 / 2,
+    ),
+    "SE_L2_Large_Lyapunov": OrbitIC(
+        name="SE_L2_Large_Lyapunov",
+        x0=1.0437961011210057,
+        z0=0,
+        vy0=-8.2206717351832326e-2,
+        half_period=8.3852287199494260 / 2,
+    ),
+    "SE_L2_Southern_Halo": OrbitIC(
+        name="SE_L2_Southern_Halo",
+        x0=1.0057032659079346,
+        z0=-1.2409260038177099e-2,
+        vy0=-1.2623398608372105e-2,
+        half_period=2.1172370586539060 / 2,
+    ),
+    "SE_Butterfly": OrbitIC(
+        name="SE_Butterfly",
+        x0=9.9478400099057562e-1,
+        z0=9.3307317319370944e-3,
+        vy0=-2.7018843432500038e-3,
+        half_period=3.6918732624381563 / 2,
+    ),
+    "SE_Dragonfly": OrbitIC(
+        name="SE_Dragonfly",
+        x0=1.0087123143716323,
+        z0=1.0173103235716078e-2,
+        vy0=-1.6680589951751042e-2,
+        half_period=5.6228423897645730 / 2,
+    ),
+}

orbit_gen/setup/system_config.py

@@ -0,0 +1,54 @@
+from dataclasses import dataclass
+import astropy.units as u
+
+
+@dataclass(frozen=True)
+class CanonicalUnits:
+    """Characteristic length/time scales for a specific CR3BP system."""
+
+    dist_km: float  # canonical distance [km]
+    time_days: float | None = None  # canonical time / (2*pi) [days]
+
+    @property
+    def du_m(self) -> float:
+        return (self.dist_km * u.km).to("m").value
+
+
+@dataclass
+class SystemConfig:
+    """Determines parameters based on the CR3BP system
+
+    Attributes:
+        name (str):
+            Label for the system
+        primary (str):
+            Primary body in the CR3BP
+        secondary (str):
+            Secondary body in the CR3BP
+
+    .. note::
+
+        "SE" for Sun/Earth "EM" for Earth/Moon, "Sun" for sun, "Earth" for Earth
+        "Moon" for moon
+
+    """
+
+    name: str
+    primary: str
+    secondary: str
+    canonical_units: CanonicalUnits
+
+
+EM_SYSTEM = SystemConfig(
+    name="EM",
+    primary="Earth",
+    secondary="Moon",
+    canonical_units=CanonicalUnits(dist_km=3.844e5),
+)
+
+SE_SYSTEM = SystemConfig(
+    name="SE",
+    primary="Sun",
+    secondary="Earth",
+    canonical_units=CanonicalUnits(dist_km=u.au.to("km")),
+)

orbit_gen/utils/eom.py

@@ -0,0 +1,265 @@
+import numpy as np
+from scipy.integrate import solve_ivp
+from scipy.optimize import fsolve
+
+
+def crtbp_eom(t, w, mu_star):
+    """Equations of motion for the CR3BP in the rotating frame with optional STM propagation
+
+    Computes the time derivative of the state in the circular restricted three-body problem
+    (CR3BP). Supports both pure state propagation and augmented state propagation including
+    a flattened 6×6 state transition matrix (STM).
+
+    Args:
+        t (float):
+            Time variable (required by ODE integrators such as `scipy.integrate.solve_ivp`
+            but not explicitly used in the CR3BP formulation).
+        w (array_like):
+            State vector. Two supported formats:
+                - len == 6: [x, y, z, vx, vy, vz]
+                - len == 42: augmented state including flattened 6×6 STM.
+        mu_star (float):
+            CR3BP mass parameter (normalized secondary mass).
+
+    Returns:
+        dw (ndarray):
+            Time derivative of the input state vector (and STM if present),
+            matching the shape of `w`.
+    """
+
+    [x, y, z, vx, vy, vz] = w[:6]
+
+    m1 = 1 - mu_star
+    m2 = mu_star
+
+    r1 = mu_star
+    r2 = 1 - mu_star
+    r_1O = r1 * np.array([-1, 0, 0])
+    r_2O = r2 * np.array([1, 0, 0])
+
+    r_PO = np.array([x, y, z])
+    v_PO = np.array([vx, vy, vz])
+
+    r_P1 = r_PO - r_1O
+    r_P2 = r_PO - r_2O
+    r1_mag = np.linalg.norm(r_P1)
+    r2_mag = np.linalg.norm(r_P2)
+
+    F_g = -m1 / r1_mag**3 * r_P1 - m2 / r2_mag**3 * r_P2
+
+    e3 = np.array([0, 0, 1])
+    a_PO = F_g - 2 * np.cross(e3, v_PO) - np.cross(e3, np.cross(e3, r_PO))
+    ax, ay, az = a_PO
+
+    if len(w) > 7:
+        dxdx = (
+            1
+            - m1 / r1_mag**3
+            - m2 / r2_mag**3
+            + 3 * m1 * (x + m2) ** 2 / r1_mag**5
+            + 3 * m2 * (x - 1 + m2) ** 2 / r2_mag**5
+        )
+        dxdy = 3 * m1 * (x + m2) * y / r1_mag**5 + 3 * m2 * (x - 1 + m2) * y / r2_mag**5
+        dxdz = 3 * m1 * (x + m2) * z / r1_mag**5 + 3 * m2 * (x - 1 + m2) * z / r2_mag**5
+        dydy = (
+            1
+            - m1 / r1_mag**3
+            - m2 / r2_mag**3
+            + 3 * m1 * y**2 / r1_mag**5
+            + 3 * m2 * y**2 / r2_mag**5
+        )
+        dydz = 3 * m1 * y * z / r1_mag**5 + 3 * m2 * y * z / r2_mag**5
+        dzdz = (
+            -m1 / r1_mag**3
+            - m2 / r2_mag**3
+            + 3 * m1 * z**2 / r1_mag**5
+            + 3 * m2 * z**2 / r2_mag**5
+        )
+
+        Phi = np.reshape(w[6:], (6, 6))
+        Z = np.zeros([3, 3])
+        I = np.identity(3)
+        A = np.array([[dxdx, dxdy, dxdz], [dxdy, dydy, dydz], [dxdz, dydz, dzdz]])
+        Omega = np.array([[0, 2, 0], [-2, 0, 0], [0, 0, 0]])
+        J = np.block([[Z, I], [A, Omega]])
+        dPhi = np.reshape(J @ Phi, (1, 36))[0]
+        dw = np.append([vx, vy, vz, ax, ay, az], dPhi)
+    else:
+        dw = [vx, vy, vz, ax, ay, az]
+
+    return dw
+
+
+def propagate(free_var, mu_star):
+    """Integrate the CR3BP equations of motion over one full orbital period
+
+    Performs numerical propagation of the CR3BP dynamics using the provided
+    initial free variables. Optionally propagates the state transition matrix
+    (STM) if included in the input.
+
+    Args:
+        free_var (array_like):
+            Initial condition vector. Supported formats:
+                - [x0, z0, vy0, T]
+                - [x0, z0, vy0, T, STM_flat] where STM_flat has length 36
+                (flattened 6×6 STM, total state dimension = 42).
+        mu_star (float):
+            CR3BP mass parameter.
+
+    Returns:
+        states (ndarray):
+            Trajectory array. Shape is either (N, 6) for state-only propagation
+            or (N, 42) when STM propagation is included.
+        times (ndarray):
+            Time vector of shape (N,) corresponding to integration output.
+    """
+
+    if len(free_var) > 4:
+        x0 = np.append(
+            [free_var[0], 0, free_var[1], 0, free_var[2], 0], free_var[4:]
+        ).tolist()
+    else:
+        x0 = [free_var[0], 0, free_var[1], 0, free_var[2], 0]
+
+    T = free_var[3]
+
+    sol = solve_ivp(crtbp_eom, [0, T], x0, args=(mu_star,), rtol=1e-12, atol=1e-12)
+    return sol.y.T, sol.t
+
+
+def constraint(free_var, mu_star):
+    """Evaluate periodicity constraints for CR3BP symmetric periodic orbit
+
+    Computes the constraint vector enforcing periodicity at the half-period
+    crossing and returns the associated state transition matrix (STM) evaluated
+    at T/2.
+
+    Args:
+        free_var (array_like):
+            Free-variable vector of the form [x0, z0, vy0, T/2] optionally
+            followed by a flattened 6×6 state transition matrix (STM).
+        mu_star (float):
+            CR3BP mass parameter (normalized secondary mass).
+
+    Returns:
+        Fx (ndarray, shape (3,)):
+            Periodicity constraint residuals evaluated at T/2:
+            typically enforcing symmetry conditions such as [y, vx, vz].
+        state_half_T (ndarray, shape (6,)):
+            State evaluated at T/2.
+        Phi (ndarray, shape (6, 6)):
+            State transition matrix evaluated at T/2.
+    """
+
+    states, _ = propagate(free_var, mu_star)
+    state = states[-1]
+
+    Fx = np.array([state[1], state[3], state[5]])
+
+    if len(state) > 7:
+        Phi = np.reshape(state[6:], (6, 6))
+    else:
+        Phi = None
+
+    state_half_T = state[0:6]
+
+    return Fx, state_half_T, Phi
+
+
+def jacobian(free_var, mu_star, state_half_T, Phi):
+    """Compute Jacobian of periodicity constraints for differential correction
+
+    Evaluates the sensitivity matrix dFx/dX for the CR3BP periodicity constraints
+    used in the single-shooting differential corrector.
+
+    Args:
+        free_var (array_like):
+            Free-variable vector [x0, z0, vy0, T/2].
+        mu_star (float):
+            CR3BP mass parameter (normalized secondary mass).
+        state_half_T (ndarray, shape (6,)):
+            State evaluated at T/2.
+        Phi (ndarray, shape (6, 6)):
+            State transition matrix evaluated at T/2.
+
+    Returns:
+        dFx (ndarray, shape (3, 4)):
+            Jacobian matrix of constraint residuals with respect to free variables.
+    """
+
+    phis = np.array(
+        [
+            [Phi[1, 0], Phi[1, 2], Phi[1, 4]],
+            [Phi[3, 0], Phi[3, 2], Phi[3, 4]],
+            [Phi[5, 0], Phi[5, 2], Phi[5, 4]],
+        ]
+    )
+
+    dw = crtbp_eom(free_var[-1], state_half_T, mu_star)
+    ddT = np.array([dw[1], dw[3], dw[5]])
+
+    dFx = np.zeros([3, 4])
+    dFx[:, 0:3] = phis
+    dFx[:, 3] = ddT
+
+    return dFx
+
+
+def continuation_constraint(w, z, sol_predictor, mu_star):
+    """System of equations for pseudo-arclength continuation (for root finding)
+
+    Defines the augmented nonlinear system used in pseudo-arclength continuation,
+    typically solved using `fsolve`. Combines the CR3BP periodicity constraints
+    with the arclength constraint relative to a predictor step.
+
+    Args:
+        w (array_like):
+            Current free-variable vector [x0, z0, vy0, T/2].
+        z (ndarray):
+            Tangent direction vector at the current solution point.
+        sol_predictor (ndarray):
+            Predicted solution from the continuation predictor step.
+        mu_star (float):
+            CR3BP mass parameter (normalized secondary mass).
+
+    Returns:
+        ndarray, shape (4,):
+            Residual vector combining periodicity constraints and arclength
+            condition, suitable for nonlinear root solving.
+    """
+
+    Fx, _, _ = constraint(w, mu_star)
+    zeq = z.T @ (w - sol_predictor)
+    return np.append(Fx, zeq)
+
+
+def jacobi_constant(pos, vel, mu_star):
+    """Compute the Jacobi constant in the circular restricted three-body problem
+
+    Evaluates the Jacobi integral in the rotating frame for a given state. The
+    Jacobi constant is conserved along CR3BP trajectories and is commonly used
+    as an invariant for validation and classification of orbits.
+
+    Args:
+        pos (array_like, shape (3,)):
+            Position vector [x, y, z] in canonical rotating-frame coordinates.
+        vel (array_like, shape (3,)):
+            Velocity vector [vx, vy, vz] in canonical units.
+        mu_star (float):
+            CR3BP mass parameter (normalized secondary mass).
+
+    Returns:
+        float:
+            Value of the Jacobi constant for the given state.
+    """
+
+    r_Mbary = np.array([1 - mu_star, 0, 0])
+    r_Ebary = np.array([-mu_star, 0, 0])
+
+    KE = np.dot(vel, vel) / 2
+    U1 = -(pos[0] ** 2 + pos[1] ** 2) / 2
+    U2 = -(
+        (1 - mu_star) / np.linalg.norm(pos - r_Ebary)
+        + mu_star / np.linalg.norm(pos - r_Mbary)
+    )
+    return KE + U1 + U2

orbit_gen/utils/units.py

@@ -0,0 +1,27 @@
+from ..setup.system_config import CanonicalUnits
+import astropy.units as u
+
+
+def time_to_canonical(dim_time, cu: CanonicalUnits):
+    return dim_time.to("day").value / cu.time_days
+
+
+def time_to_dimensional(canonical_time, cu: CanonicalUnits):
+    return canonical_time * cu.time_days * u.day
+
+
+def pos_to_canonical(dim_pos, cu: CanonicalUnits):
+    return dim_pos.to("m").value / cu.du_m
+
+
+def pos_to_dimensional(canonical_pos, cu: CanonicalUnits):
+    return canonical_pos * cu.du_m * u.m
+
+
+def vel_to_canonical(dim_vel, cu: CanonicalUnits):
+    dim_vel_mps = dim_vel.to("m/d").value
+    return (dim_vel_mps / cu.du_m) * cu.time_days
+
+
+def vel_to_dimensional(canonical_vel, cu: CanonicalUnits):
+    return canonical_vel * cu.du_m / cu.time_days * (u.m / u.day)

orbit_gen-0.1.0.dist-info/METADATA

@@ -0,0 +1,134 @@
+Metadata-Version: 2.4
+Name: orbit-gen
+Version: 0.1.0
+Summary: Tools for generating families and orbits in the CR3BP
+Author: Jayden Ripic
+Author-email: Jayden Ripic <jmr595@cornell.edu>
+License-Expression: MIT
+License-File: LICENSE
+Requires-Dist: astropy>=7.2.0
+Requires-Dist: matplotlib>=3.10.9
+Requires-Dist: numpy>=2.4.6
+Requires-Dist: platformdirs>=4.10.0
+Requires-Dist: pyyaml>=6.0.3
+Requires-Dist: scipy>=1.17.1
+Requires-Dist: spiceypy>=8.1.0
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+
+# orbit-gen
+
+Periodic orbit family generation and continuation in the Circular Restricted
+Three-Body Problem (CR3BP).
+
+## Overview
+
+`orbit-gen` generates families of periodic orbits in the CR3BP using
+single-shooting differential correction and pseudo-arclength continuation.
+It currently supports various families in the Earth-Moon and Sun-Earth systems.
+
+## Installation
+
+```bash
+pip install orbit-gen
+```
+Requires Python 3.12+.
+
+## Usage
+
+### From a config file
+
+```yaml
+# orbit_config.yaml
+system: SE                      # "EM" (Earth-Moon) or "SE" (Sun-Earth)
+preset: SE_L2_Northern_Halo
+
+max_iter: 1000                  # Maximum Newton iterations for a single differential correct attempt
+eps: 1.0e-6                     # Differential-correction convergence tolerance on residual norm
+step: 0.01                      # Pseudo-arclength continuation step size in canonical units
+max_solutions: 20               # Maximum number of orbits to generate
+max_period_days: 1000.0         # Maximum full orbit period in days
+
+output_dir: results
+show_plots: true                # Optional; for user use only
+```
+
+```python
+from orbit_gen import Family, RunConfig
+
+cfg = RunConfig.from_yaml("orbit_config.yaml")
+fam = Family(cfg)
+
+fam.save()
+if cfg.show_plots:
+    fam.plot()
+```
+
+> **Note:** constructing `Family(cfg)` runs the continuation, but does not save or plot results on its own — call `.save()` / `.plot()` explicitly, as shown above.
+
+### Programmatically
+
+```python
+from orbit_gen import Family, RunConfig
+from orbit_gen.setup.system_config import SE_SYSTEM
+from orbit_gen.setup.presets import PRESETS
+
+cfg = RunConfig(system=SE_SYSTEM, ic=PRESETS["SE_L2_Northern_Halo"])
+fam = Family(cfg)
+fam.save()
+```
+
+Available preset keys (see `orbit_gen.setup.presets.PRESETS`):
+`EM_L2_Northern_Halo`, `EM_DRO_1`, `EM_DRO_2`, `EM_L1`, `EM_Butterfly`,
+`EM_Unnamed_1`, `SE_L2_Northern_Halo`, `SE_L1_Vertical`, `SE_L2_Lyapunov`,
+`SE_L2_Small_Lyapunov`, `SE_L2_Large_Lyapunov`, `SE_L2_Southern_Halo`,
+`SE_Butterfly`, `SE_Dragonfly`.
+
+## Supported Systems
+
+| Key | System |
+|-----|--------|
+| `EM` | Earth-Moon |
+| `SE` | Sun-Earth |
+
+## Output
+
+`fam.save()` writes a `.npz` file (named after the preset, in `cfg.output_dir`)
+containing initial states, periods, and `mu_star` for all converged orbits:
+
+```python
+import numpy as np
+data = np.load("results/SE_L2_Northern_Halo.npz")
+data["states"]              # (N, 6) containing initial [x, y, z, vx, vy, vz] in canonical units for each orbit
+data["periods"]             # (N,) containing full period in canonical units for each orbit
+data["mu_star"]             # System mass ratio
+```
+
+## Contributing
+
+Clone or download the repository and run
+
+```bash
+uv sync
+```
+
+to install the package in editable mode plus dev dependencies. Run the test suite with:
+
+```bash
+uv run pytest
+```
+
+The repo includes a driver script, `scripts/crtbpModel.py`,
+which reads `orbit_config.yaml` from the project root, runs continuation,
+saves results, and shows plots.
+
+An example configuration file (`orbit_config.example.yaml`), is included in the repository, and `orbit_config.yaml` can be created and used by running:
+
+```bash
+cp orbit_config.example.yaml orbit_config.yaml   # edit as needed
+uv run python scripts/crtbpModel.py
+```
+
+## License
+
+MIT — see [LICENSE](LICENSE).

orbit_gen-0.1.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+orbit_gen/__init__.py,sha256=WwscDhNcpEt4Du2Km-jRjQHIdx2NlXqhE7wGaz_-ZKs,115
+orbit_gen/family.py,sha256=_FCp6Mk9tR2VZ8jf4mVBEf2gru6ylWmNlX0iwmkAwmo,10052
+orbit_gen/orbit.py,sha256=eaOgC3oMFKWTpKynpxQeOPETj1bsiSqhVI51y0WP2mg,4244
+orbit_gen/setup/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+orbit_gen/setup/config.py,sha256=VWNFEZqbzcouLLTMQaidEC35KDlPFTOPI-4GGT1849Y,3482
+orbit_gen/setup/kernels.py,sha256=vv7oOnD-C50SvdUDoQDR9vSCMcUnmy8t1romdhBXWU8,1497
+orbit_gen/setup/orbit_ic.py,sha256=c1Vq4DuJESXYSiP-oQf-GS4PQk4zsom8SxxuO6sxXFY,1159
+orbit_gen/setup/presets.py,sha256=TjHxsmqIMekbkNqTocWAAfU3bbi_scsCD2GCkgInIdc,2865
+orbit_gen/setup/system_config.py,sha256=5-0DdHuO6AfqsTE3K8vIIO1MEGYFzIzhiiUar-APxiw,1179
+orbit_gen/utils/eom.py,sha256=LNUCAojl1HbnRGAP3QqXtym5zV76rUUBymNIfUiNatk,8609
+orbit_gen/utils/units.py,sha256=UC84KnM5ixz_-Cnz-FKUBab6CYdCVPRHBnO-XinbOek,763
+orbit_gen-0.1.0.dist-info/licenses/LICENSE,sha256=gL4A20LPzIj-2WSYJhRqu2qJMDJbcWYveyQzUtswiK0,1069
+orbit_gen-0.1.0.dist-info/WHEEL,sha256=4OL6Foqnnp3xRY5wMkjgc25_i5YJC6dKsC6LPcjqEoU,80
+orbit_gen-0.1.0.dist-info/METADATA,sha256=L_sn__BgYYPr8tx8sXUwLBwgwYyPLcvle4m9FDyi_dA,3826
+orbit_gen-0.1.0.dist-info/RECORD,,

orbit_gen-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: uv 0.12.5
+Root-Is-Purelib: true
+Tag: py3-none-any

orbit_gen-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Jayden Ripic
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.