Progenax: Comprehensive Code, Architecture & Science Review

Document Type: Full-Package Technical Review Date: 2025-12-07 Reviewer: Claude Opus 4.5 Package Version: 0.1.0 Overall Grade: A (95/100)

Table of Contents¶

1. Executive Summary

2. Pedagogical Guide to Capabilities

3. Architecture Deep-Dive

4. Scientific Correctness Verification

5. Code Quality Assessment

6. Module-by-Module Analysis

7. Recommendations

8. Appendices

1. Executive Summary¶

1.1 Package Overview¶

progenax is a JAX-native library for generating differentiable initial conditions (ICs) for N-body gravitational simulations. It provides scientifically rigorous implementations of stellar cluster density profiles, velocity distribution functions, initial mass functions, binary orbital mechanics, and specialized transforms (mass segregation, fractal substructure, tidal truncation).

1.2 Key Metrics¶

1.3 Key Capabilities¶

1. Density Profiles: Plummer, King, Elson-Fall-Freeman (EFF) spherical models

2. Velocity DFs: Matching distribution functions with isotropic/anisotropic options

3. Initial Mass Functions: 9 models including Kroupa, Chabrier, IGIMF

4. Binary Populations: Full Keplerian orbital mechanics + period/eccentricity distributions

5. Mass Segregation: Simple radial + Baumgardt energy-ranked assignment

6. Fractal Substructure: Goodwin-Whitworth + McLuster radial overlay

7. Tidal Physics: Jacobi radius calculation and truncation

8. Two-Component Clusters: Separate populations with distinct profiles/kinematics

1.4 Design Philosophy¶

• JAX-Native: 100% JAX implementations (no numpy/scipy in core code)

• Differentiable: Inverse CDF sampling enables jax.grad() through IC generation

• Composable: Protocol-based polymorphism allows mixing any profile + velocity DF

• Immutable: Equinox modules ensure thread-safe, JIT-compatible state

• Well-Documented: Physics references throughout (30+ academic citations)

1.5 Grading Summary¶

2. Pedagogical Guide to Capabilities¶

2.1 Core Concepts¶

2.1.1 Protocol-Based Polymorphism¶

progenax uses Python’s typing.Protocol with @runtime_checkable to define three core interfaces:

from typing import Protocol, runtime_checkable
from jaxtyping import Array, Float, PRNGKeyArray

@runtime_checkable
class SpatialProfile(Protocol):
    """Interface for 3D density profiles."""

    def sample_positions(
        self,
        masses: Float[Array, "N"],
        key: PRNGKeyArray,
    ) -> Float[Array, "N 3"]:
        """Sample N particle positions from density profile."""
        ...

    def characteristic_radius(self) -> Float[Array, ""]:
        """Return characteristic scale (e.g., half-mass radius)."""
        ...

@runtime_checkable
class VelocityDF(Protocol):
    """Interface for velocity distribution functions."""

    def sample_velocities(
        self,
        positions: Float[Array, "N 3"],
        masses: Float[Array, "N"],
        key: PRNGKeyArray,
        G: float | None = None,  # Uses jaxstro.units.DEFAULT.G if None
    ) -> Float[Array, "N 3"]:
        """Sample velocities at given positions."""
        ...

@runtime_checkable
class IMFProtocol(Protocol):
    """Interface for initial mass functions."""
    m_min: float
    m_max: float

    def sample(self, key: PRNGKeyArray, n: int) -> Float[Array, "n"]:
        """Sample n stellar masses."""
        ...

    def ppf(self, u: Float[Array, "..."]) -> Float[Array, "..."]:
        """Percent-point function (inverse CDF)."""
        ...

Why Protocols? They enable compositional IC assembly:

# Mix any profile with any velocity DF!
from jaxstro.units import STELLAR
from progenax import PlummerProfile, KingVelocityDF, build_spatial_ic

profile = PlummerProfile(r_h=1.0)    # Plummer density
df = KingVelocityDF(r_t=2.0)         # King velocities
ic = build_spatial_ic(profile, masses, df, key, G=STELLAR.G)  # ~0.00450

2.1.2 The ICResult Workflow¶

The central output is ICResult, a frozen dataclass:

@dataclass(frozen=True)
class ICResult:
    positions: Float[Array, "N 3"]      # Cartesian [length units]
    velocities: Float[Array, "N 3"]     # Cartesian [velocity units]
    masses: Float[Array, "N"]           # [M_sun]
    softening: float                    # [length units]
    stellar_radii: Float[Array, "N"]    # [R_sun]
    ids: Optional[Float[Array, "N"]]    # Particle IDs

The 7-step IC generation workflow (build_spatial_ic()):

1. Split RNG key → separate keys for positions, velocities

2. Sample positions → from SpatialProfile

3. Sample velocities → from VelocityDF (position-dependent)

4. Compute softening → ε = 0.01 × r_char / N^(1/3)

5. Compute stellar radii → mass-radius relation (MS + brown dwarfs)

6. Transform to COM frame → zero total momentum

7. Apply virial scaling → scale velocities to target Q = 2T/|V|

2.1.3 Equinox Immutability Pattern¶

All stateful classes use Equinox modules:

import equinox as eqx
import jax.numpy as jnp
from jaxtyping import Array, Float

class PlummerProfile(eqx.Module):
    """Immutable Plummer profile as Equinox PyTree."""

    r_h: Float[Array, ""]  # Half-mass radius
    a: Float[Array, ""]    # Scale radius (computed)

    def __init__(self, r_h: float = 1.0):
        self.r_h = jnp.asarray(r_h, dtype=jnp.float64)
        self.a = self.r_h * jnp.sqrt((1.0 - 0.5**(2/3)) / 0.5**(2/3))

Benefits:

• Immutable (hashable, safe to pass through JIT)

• PyTree-compatible (works with jax.tree_util)

• Clean parameter handling

2.2 Module Catalog¶

2.3 Usage Examples¶

Example 1: Basic Plummer Cluster IC¶

import jax
import jax.numpy as jnp
from progenax import (
    PlummerProfile,
    PlummerVelocityDF,
    build_spatial_ic,
)
from progenax.imf import PowerLawIMF

# 1. Sample masses from Kroupa IMF
key = jax.random.PRNGKey(42)
key_imf, key_ic = jax.random.split(key)

imf = PowerLawIMF.kroupa()
masses = imf.sample(key_imf, 1000)  # 1000 stars

# 2. Create profile and velocity DF
profile = PlummerProfile(r_h=1.0)        # 1 pc half-mass radius
velocity_df = PlummerVelocityDF(r_h=1.0)

# 3. Generate ICs
G = 0.00450  # pc³ M☉⁻¹ Myr⁻²
ic = build_spatial_ic(
    profile=profile,
    masses=masses,
    velocity_df=velocity_df,
    key=key_ic,
    G=G,
    Q=1.0,  # Virial equilibrium
)

print(f"Positions shape: {ic.positions.shape}")   # (1000, 3)
print(f"Velocities shape: {ic.velocities.shape}") # (1000, 3)
print(f"Softening: {ic.softening:.4f} pc")

Example 2: Two-Component Cluster with Mass Segregation¶

from progenax import (
    PlummerProfile,
    PlummerVelocityDF,
    TwoComponentConfig,
    generate_two_component_cluster,
)
from progenax.profiles import apply_mass_segregation_baumgardt

# Configure two-component cluster
config = TwoComponentConfig(
    f_A=0.3,                              # 30% in extended halo
    profile_A=PlummerProfile(r_h=2.0),    # Extended population
    profile_B=PlummerProfile(r_h=0.5),    # Concentrated population
    velocity_df_A=PlummerVelocityDF(r_h=2.0),
    velocity_df_B=PlummerVelocityDF(r_h=0.5),
)

# Generate cluster
positions, velocities, pop_id = generate_two_component_cluster(
    masses, config, key, G=0.00450
)

# Apply Baumgardt-style mass segregation (s=0.8 = strong)
positions_seg, velocities_seg = apply_mass_segregation_baumgardt(
    positions, velocities, masses, s=0.8, key=key, G=0.00450
)

Example 3: Binary Population with Radial Fraction¶

from progenax.binaries import (
    RadialBinaryFraction,
    MassDependentBinaryConfig,
    LogNormalPeriod,
    SanaOBPeriod,
    ThermalEccentricity,
    MoeEccentricity,
    sample_mass_dependent_orbits,
)

# Radial binary fraction: more binaries in core
rbf = RadialBinaryFraction(
    fb0=0.5,      # Baseline 50%
    A=0.5,        # Core enhancement
    alpha=1.0,    # Power-law index
    r_scale=1.0,  # Scale radius
)

# Sample binary membership
radii = jnp.linalg.norm(positions, axis=1)
is_binary = rbf.sample_membership(radii, key)

# Mass-dependent orbital parameters
config = MassDependentBinaryConfig(
    m_break=8.0,  # M☉ threshold
    low_mass_period=LogNormalPeriod(mu_log_P=4.8, sigma_log_P=2.3),
    high_mass_period=SanaOBPeriod(),  # O/B stars: shorter periods
    low_mass_eccentricity=ThermalEccentricity(),
    high_mass_eccentricity=MoeEccentricity(),  # Period-dependent
)

periods, eccentricities = sample_mass_dependent_orbits(masses, config, key)

3. Architecture Deep-Dive¶

3.1 Design Patterns¶

Pattern 1: Protocol-Based Polymorphism¶

Location: protocols.py Purpose: Enable compositional IC assembly

# Runtime checking enables duck typing with verification
from typing import Protocol, runtime_checkable

@runtime_checkable
class SpatialProfile(Protocol):
    ...

# Usage: verify any object implements protocol
profile = MyCustomProfile(...)
assert isinstance(profile, SpatialProfile)  # Runtime check!

Implementations:

• SpatialProfile: PlummerProfile, KingProfile, EFFProfile

• VelocityDF: PlummerVelocityDF, KingVelocityDF, EFFVelocityDF

• IMFProtocol: PowerLawIMF, ChabrierIMF, IGIMF, ...

Pattern 2: Equinox Module Immutability¶

Location: All stateful classes Purpose: PyTree compatibility, thread safety, JIT compilation

import equinox as eqx

class PlummerProfile(eqx.Module):
    r_h: Float[Array, ""]
    a: Float[Array, ""]

    # Immutable: no __setattr__ after init
    # PyTree: compatible with jax.tree_util
    # Hashable: can be used as dict keys

Pattern 3: Explicit G Parameter¶

Location: All physics functions Purpose: Unit system flexibility, testability

# CORRECT: Explicit G parameter
def compute_potential_energy(positions, masses, G, softening=0.0):
    ...

# WRONG: Global state
def compute_potential_energy(positions, masses, softening=0.0):
    G = get_G()  # Hidden dependency!
    ...

Note: Unlike gravax (which uses get_G()), progenax requires explicit G throughout.

Pattern 4: Inverse CDF Sampling (Differentiable)¶

Location: PlummerProfile._sample_radii(), PlummerVelocityDF._sample_velocity_magnitudes() Purpose: Enable jax.grad() through IC generation

# Plummer radii via inverse CDF
u = jax.random.uniform(key, (N,))      # u ~ U(0,1)
u_23 = jnp.power(u, 2.0/3.0)
radii = a * jnp.sqrt(u_23 / (1.0 - u_23))  # Inverse CDF

# Plummer velocities via Beta distribution
u = jax.random.beta(key, a=1.5, b=4.5, shape=(N,))  # u = q²
v = jnp.sqrt(u) * v_esc  # EXACT (no rejection sampling!)

Why This Matters: Rejection sampling blocks gradients; inverse CDF is differentiable.

Pattern 5: Fixed-Iteration Loops via jax.lax.scan¶

Location: apply_mass_segregation_baumgardt(), Kepler solver Purpose: JIT compatibility, differentiability

# CORRECT: Fixed iterations via scan
def assign_step(carry, i):
    available_mask, assignments, key = carry
    # ... assignment logic ...
    return (new_mask, new_assignments, key), None

(final_mask, assignments, _), _ = jax.lax.scan(
    assign_step,
    (initial_mask, initial_assignments, key),
    jnp.arange(N)  # Fixed N iterations
)

# WRONG: Convergence loop (not differentiable)
while error > tolerance:  # Variable iterations
    ...

3.2 Critical Files¶

__init__.py (117 LOC)¶

Central namespace with 57 public exports organized by category:

• Protocols (3)

• Spatial Profiles (4)

• Velocity DFs (3)

• IC Builders (7)

• Analytical Test Cases (15)

• Binary Mechanics (5)

• Tidal Physics (4)

• Fractal Substructure (3)

• Two-Component Populations (2)

protocols.py (120 LOC)¶

Three runtime-checkable protocols:

1. SpatialProfile: sample_positions(), characteristic_radius()

2. VelocityDF: sample_velocities()

3. IMFProtocol: logpdf(), cdf(), ppf(), sample(), mean_mass()

Issue: VelocityDF protocol doesn’t include G parameter, but all implementations require it.

builders.py (289 LOC)¶

Core IC assembly utilities:

• ICResult: Frozen output dataclass

• compute_stellar_radii(): Mass-radius relation (3 regimes)

• compute_kinetic_energy(): T = 0.5 × Σ m_i v_i²

• compute_potential_energy(): V = -G × Σ_{i<j} m_i m_j / r_ij

• to_com_frame(): Center-of-mass transformation

• virial_scale(): Scale velocities to target Q = 2T/|V|

• build_spatial_ic(): Master 7-step workflow

3.3 Dependencies¶

# Core dependencies
jax >= 0.4.20
jaxlib >= 0.4.20
equinox >= 0.11.0
jaxtyping >= 0.2.25
diffrax >= 0.4.0  # For King profile ODE solver

# Optional
jaxstro  # Shared utilities (unit systems, coordinates)

# Dev dependencies
pytest >= 7.0
pytest-cov

Critical: No numpy/scipy in core code. All operations use JAX primitives.

3.4 Package Structure¶

progenax/
├── src/progenax/
│   ├── __init__.py          # Public API (57 exports)
│   ├── protocols.py          # 3 protocols
│   ├── builders.py           # ICResult + build_spatial_ic
│   ├── populations.py        # Two-component clusters
│   ├── tidal.py              # Jacobi radius + truncation
│   ├── profiles/
│   │   ├── __init__.py
│   │   ├── plummer.py        # PlummerProfile
│   │   ├── king.py           # KingProfile + ODE solver
│   │   ├── eff.py            # EFFProfile
│   │   └── mass_segregation.py
│   ├── kinematics/
│   │   ├── __init__.py
│   │   ├── plummer_df.py     # PlummerVelocityDF (Beta sampling)
│   │   ├── king_df.py        # KingVelocityDF
│   │   ├── eff_df.py         # EFFVelocityDF
│   │   ├── anisotropy.py     # Osipkov-Merritt
│   │   └── rotation.py       # Solid body, differential
│   ├── imf/
│   │   ├── __init__.py
│   │   ├── base.py           # BaseIMF
│   │   ├── power_law.py      # PowerLawIMF (Kroupa, Salpeter)
│   │   ├── chabrier.py       # ChabrierIMF
│   │   ├── environment.py    # EnvironmentIMF
│   │   ├── binary.py         # BinaryIMF
│   │   └── igimf.py          # IGIMF
│   ├── binaries/
│   │   ├── __init__.py
│   │   ├── kepler.py         # KeplerElements, Kepler solver
│   │   ├── orbital_state.py  # BinaryOrbitalState
│   │   └── population.py     # Period/eccentricity distributions
│   ├── substructure/
│   │   ├── __init__.py
│   │   └── fractal.py        # Goodwin-Whitworth + overlays
│   └── analytical/
│       ├── __init__.py
│       └── core.py           # Solar system, Kepler orbits
├── tests/
│   ├── unit/                 # ~300 unit tests
│   └── integration/          # ~50 integration tests
└── docs/
    └── core-reviews/         # This document

4. Scientific Correctness Verification¶

4.1 Plummer Profile (Plummer 1911)¶

Density:

Scale radius from half-mass radius:

Inverse CDF for radii:

4.2 Plummer Velocity DF (Dehnen 1993)¶

Distribution function:

Velocity magnitude distribution (q = v/v_esc):

Key insight: Let u = q², then u ~ Beta(3/2, 9/2). This is EXACT (no rejection sampling!).

Escape velocity:

Velocity dispersion relation:

4.3 King Profile (King 1966)¶

Dimensionless Poisson equation:

K-function:

Density relation:

4.4 Mass Segregation (Baumgardt+2008)¶

Energy-ranked orbit assignment algorithm:

1. Compute binding energies: E_k = ½v² + Φ(r)

2. Sort orbits by energy (most bound first)

3. Sort masses descending (most massive first)

4. Assign masses to orbits via power-law slot selection:

   • j = ⌊(n_remaining - 1) × (1 - U^(1-s))⌋

   • s = 0: random, s = 1: maximal segregation

Note: Uses approximate potential Φ ≈ -GM_total/r (documented trade-off).

4.5 Fractal Substructure (Goodwin & Whitworth 2004)¶

Survival probability:

where D ∈ [1.5, 3.0] is the fractal dimension:

• D = 1.5: p = 0.35 (very clumpy)

• D = 2.0: p = 0.5 (moderately clumpy)

• D = 3.0: p = 1.0 (uniform)

Algorithm:

1. Root node at origin

2. Each parent spawns 8 children (cube corners)

3. Each child survives with probability p

4. Repeat for g_max generations

5. Filter to unit sphere, downsample to N

4.6 Binary Period Distributions¶

4.7 Eccentricity Distributions¶

Moe+2017 details:

• P < P_circ (~10d): Tidally circularized (e ≈ 0)

• P > P_thermal (~1000d): Thermal distribution

• Smooth logistic transition for differentiability

4.8 Jacobi Radius (King 1962)¶

Point-mass approximation:

Isothermal halo:

4.9 References Verified¶

1. Plummer (1911) MNRAS 71, 460 - Original Plummer model

2. King (1962) AJ 67, 471 - Tidal radius

3. King (1966) AJ 71, 64 - Lowered isothermal model

4. Heggie (1975) MNRAS 173, 729 - Thermal eccentricity

5. Dehnen (1993) MNRAS 265, 250 - Exact Plummer DF

6. Duquennoy & Mayor (1991) A&A 248, 485 - Solar-type binary periods

7. Chabrier (2003) PASP 115, 763 - Lognormal IMF

8. Goodwin & Whitworth (2004) A&A 413, 929 - Fractal clusters

9. Baumgardt+2008 MNRAS 384, 1231 - Primordial mass segregation

10. Kupper+2011 MNRAS 417, 2300 - McLuster

11. Sana+2012 Science 337, 444 - O-star binary fraction

12. Moe & Di Stefano (2017) ApJS 230, 15 - Binary statistics review

5. Code Quality Assessment¶

5.1 Strengths¶

Excellent Documentation¶

Every module includes:

• Purpose statement

• Mathematical formulas in docstrings

• References to academic literature

• Usage examples

• Parameter descriptions with units

Example from PlummerVelocityDF:

class PlummerVelocityDF(eqx.Module):
    """
    Plummer (1911) velocity distribution function.

    Samples velocity magnitudes from the exact Plummer DF using Beta distribution
    (no rejection sampling required). Velocities are isotropically distributed.

    The distribution for q = v/v_esc is:
        g(q) ∝ q² (1 - q²)^(7/2)  for q ∈ [0, 1]

    Sampling method:
        Let u = q², then u ~ Beta(3/2, 9/2)
        Therefore: q = sqrt(u), v = q × v_esc

    References:
        Plummer (1911), MNRAS, 71, 460 - Original Plummer model
        Dehnen (1993), MNRAS, 265, 250 - Exact analytical DF
    """

JAX-Native Throughout¶

• ✅ All array operations use jax.numpy

• ✅ No numpy or scipy in core code

• ✅ All sampling uses jax.random

• ✅ Control flow via jax.lax.cond, jax.lax.scan

• ✅ Vectorization via jax.vmap

• ✅ JIT compilation via @jax.jit

Immutable Data Structures¶

• ✅ All stateful classes are Equinox modules

• ✅ Output is frozen dataclass (ICResult)

• ✅ No in-place mutations

• ✅ Thread-safe by design

Comprehensive Test Suite¶

• 350+ tests

• 9:1 test-to-code ratio

• Tests cover:

  • Output shapes

  • Value ranges

  • Statistical properties

  • Edge cases

  • Reproducibility

  • Differentiability

5.2 Issues Found¶

Note: All “Minor” issues from the original review have been resolved. The G unit system fix (commit 6619f25) changed the VelocityDF protocol to use G: float | None = None with runtime resolution to jaxstro.units.DEFAULT.G.

5.3 Code Metrics¶

6. Module-by-Module Analysis¶

6.1 profiles/plummer.py (124 LOC)¶

Purpose: Plummer (1911) spherical density profile

Key Class: PlummerProfile(eqx.Module)

• r_h: Float[Array, ""] - Half-mass radius

• a: Float[Array, ""] - Scale radius (computed)

Methods:

• sample_positions(masses, key) → (N, 3) Cartesian positions

• characteristic_radius() → r_h scalar

Implementation Details:

• Inverse CDF sampling (differentiable)

• Isotropic angle generation

• JIT-decorated internal method

Quality: ✅ Excellent - clean, well-documented, verified formulas

6.2 profiles/king.py (~400 LOC)¶

Purpose: King (1966) lowered isothermal model

Key Class: KingProfile(eqx.Module)

• W_0: float - Central dimensionless potential (1-12)

• r_c: float - Core radius

• r_t: float - Tidal radius

Key Function: solve_king_profile(W_0, n_points=1000)

• ODE integration via diffrax.Tsit5

• Returns (xi, psi, rho_tilde) on grid

Implementation Details:

• Numerical CDF via trapezoid integration

• Linear interpolation for inverse sampling

• L’Hôpital’s rule at ξ=0 for ODE stability

Quality: ✅ Good - correct physics, could add validation against LIMEPY

6.3 profiles/mass_segregation.py (251 LOC)¶

Purpose: Primordial mass segregation transforms

Key Functions:

1. apply_mass_segregation(positions, masses, eta, m_ref) - Simple radial scaling

2. compute_mass_segregation_ratio(positions, masses, threshold) - MSR diagnostic

3. apply_mass_segregation_baumgardt(positions, velocities, masses, s, key, G) - Energy-ranked

Implementation Details:

• Baumgardt uses jax.lax.scan for slot selection (differentiable)

• Cumsum trick for finding k-th available slot

• Approximate binding energy: E = 0.5v² - GM/r

Quality: ✅ Good - correct algorithm, documented approximations

6.4 kinematics/plummer_df.py (197 LOC)¶

Purpose: Exact Plummer velocity distribution function

Key Class: PlummerVelocityDF(eqx.Module)

• r_h: Float[Array, ""] - Half-mass radius

• a: Float[Array, ""] - Scale radius

Methods:

• sample_velocities(positions, masses, key, G) → (N, 3) velocities

• _sample_velocity_magnitudes(r, masses, key, G) → (N,) magnitudes

Implementation Details:

• Beta(3/2, 9/2) sampling for velocity magnitudes (EXACT!)

• No rejection sampling required

• Isotropic direction generation

Quality: ✅ Excellent - best implementation of Plummer DF I’ve seen

6.5 kinematics/anisotropy.py (~110 LOC)¶

Purpose: Velocity anisotropy transforms

Key Function: apply_osipkov_merritt(velocities, positions, key, r_a)

Formula:

Implementation Details:

• Decompose velocity into radial/tangential components

• Scale to match target σ_r/σ_t ratio

• Preserve total speed

Quality: ✅ Good - correct physics, handles edge cases

6.6 kinematics/rotation.py (~110 LOC)¶

Purpose: Solid body and differential rotation

Key Functions:

1. apply_solid_body_rotation(velocities, positions, omega, axis) - v_rot = ω × r

2. apply_differential_rotation(velocities, positions, v_peak, R_peak, axis) - v_φ(R) profile

Differential rotation formula:

Quality: ✅ Good - clean implementation

6.7 binaries/kepler.py (~400 LOC)¶

Purpose: Keplerian orbital mechanics

Key Class: KeplerElements(eqx.Module)

• a, e, i, Omega, omega, M_0 - 6 orbital elements

• m1, m2 - Component masses

Methods:

• to_state(G) → (r, v) relative position/velocity

• to_binary_state(G) → (r1, v1, r2, v2) barycentric

• from_state(r, v, m1, m2, G) → KeplerElements (classmethod)

Key Function: _solve_kepler_equation(M, e, n_iter=10)

• Newton-Raphson via jax.lax.scan

• Fixed iterations for JIT compatibility

Quality: ⚠️ Good with issues - division protection could be tighter (1e-20 → 1e-12)

6.8 binaries/population.py (~530 LOC)¶

Purpose: Binary period/eccentricity distributions

Key Classes:

• LogUniformPeriod - Öpik’s law

• LogNormalPeriod - Duquennoy & Mayor 1991

• SanaOBPeriod - O/B star periods (Sana+2012)

• ThermalEccentricity - f(e) = 2e

• MoeEccentricity - Period-dependent (Moe+2017)

• RadialBinaryFraction - f_b(r) spatial variation

• MassDependentBinaryConfig - Routing by mass threshold

Quality: ✅ Good - comprehensive coverage, well-referenced

6.9 substructure/fractal.py (~320 LOC)¶

Purpose: Fractal substructure generation

Key Functions:

1. generate_fractal_positions(n_stars, key, d_fractal, g_max=6) - Goodwin-Whitworth

2. apply_fractal_overlay_radial(positions_smooth, key, d_fractal) - McLuster-style

3. apply_fractal_overlay_blend(positions_smooth, key, d_fractal, lambda_frac) - Linear blend

Implementation Details:

• Static allocation for JIT (8^g_max cells)

• Boolean masking for alive particles

• Weighted random sampling for downsampling

Quality: ✅ Good - correct algorithm, clearly documented trade-offs

6.10 populations.py (185 LOC)¶

Purpose: Two-component cluster generation

Key Class: TwoComponentConfig

• f_A: float - Fraction in population A

• profile_A, profile_B - Spatial profiles

• velocity_df_A, velocity_df_B - Velocity DFs

Key Function: generate_two_component_cluster(masses, config, key, G, pop_mask=None)

Implementation Details:

• Always samples both populations (JIT requirement)

• Uses jnp.where for selection

• Optional custom population mask

Quality: ✅ Good - clean design, JIT-compatible

6.11 tidal.py (150 LOC)¶

Purpose: Tidal physics

Key Functions:

• jacobi_radius(R, M_cluster, M_galaxy) - Point-mass host

• jacobi_radius_isothermal(R, M_cluster, V_circ) - SIS halo

• apply_tidal_truncation(positions, velocities, masses, r_t) - Sharp cutoff

• fill_factor_to_r_h(fill_factor, r_J) - r_h from fill factor

Quality: ✅ Good - correct formulas, could add smooth truncation option

6.12 imf/ Package (~1200 LOC total)¶

Implementations:

1. PowerLawIMF - N-segment piecewise power-law (Kroupa, Salpeter)

2. ChabrierIMF - Lognormal + power-law (Chabrier 2003)

3. EnvironmentIMF - Density/temperature dependent

4. BinaryIMF - Companion mass ratios

5. IGIMF - Integrated galactic IMF

Quality: ✅ Excellent - comprehensive coverage, correct normalization

6.13 analytical/ Package (~200 LOC)¶

Test Cases:

• two_body_kepler(a, e, m1, m2, G) - Kepler orbit

• earth_sun_2body() - Realistic Earth-Sun

• three_body_figure_eight() - Periodic 3-body (Chenciner & Montgomery)

• harmonic_oscillator(A, omega) - 1D/2D SHO

• solar_system_full() - All 8 planets + Sun

• solar_system_inner_4() - Sun + inner 4 planets

Data Source: NASA JPL Horizons (J2000.0 epoch)

Quality: ✅ Excellent - well-documented, high-precision data

7. Recommendations¶

7.1 Critical (Should Fix) — ✅ ALL COMPLETED¶

7.1.1 Add G Parameter to VelocityDF Protocol — ✅ COMPLETED¶

Location: protocols.py:56-80

Issue: Protocol signature doesn’t include G, but all implementations require it.

Resolution (commit 6619f25):

@runtime_checkable
class VelocityDF(Protocol):
    def sample_velocities(
        self,
        positions: Float[Array, "N 3"],
        masses: Float[Array, "N"],
        key: PRNGKeyArray,
        G: float | None = None,  # ✅ IMPLEMENTED
    ) -> Float[Array, "N 3"]:
        """
        Args:
            G: Gravitational constant. If None, uses jaxstro.units.DEFAULT.G
               (~0.00450 for stellar dynamics in pc³ Msun⁻¹ Myr⁻²)
        """
        ...

Implementation pattern in velocity DFs:

def sample_velocities(self, positions, masses, key, G=None):
    if G is None:
        from jaxstro.units import DEFAULT
        G = DEFAULT.G  # ~0.00450 for STELLAR, ~39.478 for PLANETARY
    # ... use G

7.1.2 Strengthen Division Protection in Kepler Solver — ✅ COMPLETED¶

Location: kepler.py:122, 125

Issue: Uses 1e-20 which is too loose for e→1 (parabolic limit).

Resolution: Changed to 1e-12 minimum denominator.

7.2 Important (Should Add)¶

7.2.1 Add RadialBinaryFraction Unit Tests¶

Location: tests/unit/binaries/test_population.py

Missing Tests:

• test_compute_shape()

• test_fb_in_range() (clipped to [0, 1])

• test_A_positive_core_enhanced()

• test_A_negative_core_depleted()

• test_sample_membership_statistics()

7.2.2 Validate King Profile Against LIMEPY¶

Purpose: Ensure ODE solution matches established code

Test:

def test_king_profile_matches_limepy():
    """Compare King profile to LIMEPY for W0 = 5, 7, 9."""
    for W0 in [5.0, 7.0, 9.0]:
        # Generate our profile
        xi, psi, rho = solve_king_profile(W0)

        # Compare to LIMEPY output
        # ... (need to add LIMEPY as test dependency)

7.3 Nice to Have¶

7.3.1 Document Non-Differentiability of Radial Overlay¶

Location: fractal.py:187

Current docstring mentions it, but could be more prominent:

def apply_fractal_overlay_radial(...):
    """...

    Warning:
        This function is NOT differentiable due to sorting operations.
        Use apply_fractal_overlay_blend() for gradient-based optimization.
    """

7.3.2 Add Smooth Tidal Truncation Option¶

Location: tidal.py

Current: Sharp cutoff (remove all r > r_t)

Enhancement: King-style smooth truncation:

def apply_smooth_tidal_truncation(
    positions, velocities, masses, r_t, r_c
):
    """Apply King-style smooth truncation."""
    # Weight particles by King profile density ratio
    ...

7.3.3 Centralize CGS Constants¶

Location: imf/environment.py:31-45

Issue: Hardcoded constants (k_B, m_p, G_cgs) should use jaxstro utilities.

Fix:

# Replace hardcoded values with:
from jaxstro.constants import k_B_cgs, m_p_cgs, G_cgs

8. Appendices¶

Appendix A: Complete Public API¶

from progenax import (
    # Protocols
    SpatialProfile,
    VelocityDF,
    IMFProtocol,

    # Spatial Profiles
    PlummerProfile,
    KingProfile,
    EFFProfile,
    solve_king_profile,

    # Velocity DFs
    PlummerVelocityDF,
    KingVelocityDF,
    EFFVelocityDF,

    # IC Builders
    ICResult,
    build_spatial_ic,
    to_com_frame,
    virial_scale,
    compute_kinetic_energy,
    compute_potential_energy,
    compute_stellar_radii,

    # Analytical Test Cases
    AnalyticalIC,
    two_body_kepler,
    earth_sun_2body,
    three_body_figure_eight,
    harmonic_oscillator,
    harmonic_oscillator_2d,
    plummer_fixed_particles,
    solar_system_full,
    solar_system_inner_4,
    # ... (15 total)

    # Binary Mechanics
    KeplerElements,
    compute_period,
    period_to_semimajor_axis,
    BinaryOrbitalState,
    batch_elements_to_resolved,

    # Tidal Physics
    jacobi_radius,
    jacobi_radius_isothermal,
    apply_tidal_truncation,
    fill_factor_to_r_h,

    # Fractal Substructure
    generate_fractal_positions,
    apply_fractal_overlay_radial,
    apply_fractal_overlay_blend,

    # Two-Component Populations
    TwoComponentConfig,
    generate_two_component_cluster,
)

Appendix B: Test Coverage by Module¶

Appendix C: Performance Characteristics¶

Notes:

• All times on Apple M1 with JIT compilation

• First call includes compilation overhead (~1-5s)

• King profile limited by ODE solver (not JIT-compiled)

• Baumgardt segregation O(N²) due to energy calculation

Appendix D: Glossary¶

Document Metadata¶

This review was generated using systematic exploration of the progenax codebase with parallel analysis of core architecture, spatial profiles/kinematics, and specialized modules. All scientific formulas were verified against published literature.