`progenax.dynamics`

progenax.dynamics¶

Auto-generated by scripts/build_api_reference.py from progenax source. Re-run after public-API changes; the script is idempotent.

Module path: progenax/dynamics/

Public symbols: 6

Contents¶

• compute_kinetic_energy

• compute_potential_energy

• compute_virial_ratio

• mass_group_masks

• per_group_virial_ratio

• rescale_velocities_to_virial

dynamics.compute_kinetic_energy¶

function

compute_kinetic_energy(velocities: Float[Array, 'N 3'], masses: Float[Array, 'N']) -> Float[Array, '']

Compute total kinetic energy: T = 0.5 * sum(m_i * v_i^2).

Source: src/progenax/dynamics/virial.py#L16

dynamics.compute_potential_energy¶

function

compute_potential_energy(positions: Float[Array, 'N 3'], masses: Float[Array, 'N'], G: float, softening: float = 0.0, block_size: int = 256) -> Float[Array, '']

Total potential energy V = -G * sum_{i<j} m_i m_j / r_ij (Plummer-softened).

Blocked row-scan (lax.scan over row blocks of block_size stars vs ALL columns): peak transient memory is O(block_size * N), not O(N^2), for the forward AND backward pass — the dense kernel measured 32.8 GB at N = 2e4 (2026-06-10); blocked at the default 256 it is ~0.12 GB. The backward pass stays O(block_size * N) via jax.checkpoint rematerialization: each block’s forward is recomputed during the vjp instead of stored, so no O(N^2) stacked residuals accumulate across scan iterations. Identical pair set and per-pair arithmetic; only float64 summation ORDER changes across blocks (re-association at the 1e-15 relative level). block_size is a Python int and must be static under jax.jit.

Differentiable at softening=0: the i<j mask feeds excluded entries (diagonal, lower triangle, padded rows) a safe value before sqrt so no masked-out sqrt(0) cotangent can NaN-poison the gradient. This is the single canonical energy implementation; progenax.builders re-exports it.

Source: src/progenax/dynamics/virial.py#L34

dynamics.compute_virial_ratio¶

function

compute_virial_ratio(positions: Float[Array, 'N 3'], velocities: Float[Array, 'N 3'], masses: Float[Array, 'N'], G: float, softening: float = 0.0) -> Float[Array, '']

Compute virial ratio Q = T / |V|.

Convention: Q = 0.5: Virial equilibrium (2T + V = 0) Q < 0.5: Subvirial (cold, collapsing) Q > 0.5: Supervirial (hot, expanding)

Args

Returns: Virial ratio Q = T / |V|

Source: src/progenax/dynamics/virial.py#L85

dynamics.mass_group_masks¶

function

mass_group_masks(masses: Float[Array, 'N'], n_groups: int) -> Bool[Array, 'n_groups N']

Partition stars into n_groups equal-count mass-rank bins (light -> heavy).

Returns a boolean array (n_groups, N) where row g selects the stars whose mass rank falls in group g (group 0 = lightest, group n_groups-1 = heaviest). The groups are disjoint and cover all N stars (up to integer division of N).

Used to ask whether each mass sub-population is individually in virial equilibrium (:func:per_group_virial_ratio) — the diagnostic that distinguishes a true multi-mass equilibrium from a globally-rescaled blend.

Source: src/progenax/dynamics/virial.py#L114

dynamics.per_group_virial_ratio¶

function

per_group_virial_ratio(positions: Float[Array, 'N 3'], velocities: Float[Array, 'N 3'], masses: Float[Array, 'N'], G: float, group_masks: Bool[Array, 'n_groups N'], softening: float = 0.0) -> Float[Array, 'n_groups']

Per-mass-group virial ratio Q_j = T_j / |W_j| in the TOTAL gravitational field.

Scalar virial theorem for a subsystem in steady state: 2 T_j + W_j = 0, where

T_j = 0.5 sum_{i in j} m_i v_i^2        (kinetic energy of the group)
W_j = sum_{i in j} m_i r_i . a_i        (a_i = acceleration from ALL stars)

so Q_j = T_j / |W_j| = 0.5 for a group in equilibrium — the same convention as :func:compute_virial_ratio. By the Clausius identity sum_i m_i r_i . a_i = V for 1/r gravity, a single all-ones group reproduces the global virial ratio exactly (W = V), and the per-group W_j partition the global virial.

Velocities are measured relative to the mass-weighted mean (bulk motion removed) so a moving COM does not inflate T_j.

.. warning:: Per-group W_j = sum_{i in j} m_i r_i . a_i is ORIGIN-DEPENDENT for a proper subgroup: under r -> r + d it shifts by d . sum_{i in j} m_i a_i, which is nonzero because a subgroup feels a net force from the rest of the cluster (only the GLOBAL sum is origin-independent, sum_all m_i a_i = 0). Callers must therefore pass positions PRE-CENTERED on the cluster COM (all in-tree callers do); an off-center origin biases the per-group Q_j (audit S16).

Args

Returns: Q_j for each group, shape (n_groups,).

Source: src/progenax/dynamics/virial.py#L179

dynamics.rescale_velocities_to_virial¶

function

rescale_velocities_to_virial(positions: Float[Array, 'N 3'], velocities: Float[Array, 'N 3'], masses: Float[Array, 'N'], G: float, target_Q: float = 0.5, softening: float = 0.0) -> Float[Array, 'N 3']

Rescale velocities to achieve target virial ratio Q = T/|V|.

This is the single implementation of virial velocity rescaling; builders.virial_scale delegates here (audit J5 dedupe).

Cold input (T=0 -> Q_current=0) makes the rescale undefined (0/0). A CONCRETE T=0 refuses loudly (ValueError); a TRACED T=0 yields NaN — the honest sentinel, since a traced scalar cannot gate a Python raise (audit J5; the previous +1e-10 hack silently mapped cold input to ~0 velocities).

Args

Returns: Rescaled velocities with Q = target_Q

Source: src/progenax/dynamics/virial.py#L249