Fit the half-mass radius with jax.grad

Goal. Recover a cluster’s half-mass radius $\rh$ by differentiating through the IC generator itself. progenax samples every IC with jax.lax.scan over a fixed iteration count (never while_loop), so build_spatial_ic is differentiable end to end: jax.grad flows from an observable all the way back to $\rh$ .

Inputs and assumptions¶

Table 1:Recipe inputs

Input	Meaning and role	Fiducial
$\rh^{\rm true}$	The half-mass radius of the synthetic “observed” cluster.	1.3 pc (recovered)
observable	The scalar matched in the loss. We use the RMS radius, which is smooth and monotone in $\rh$ .	RMS radius
`key`	Frozen PRNG key — held fixed across the optimization so the forward map is deterministic.	`PRNGKey(0)`
`G` (`STELLAR.G`)	Required gravitational constant; threaded through the velocity sampling and virial scaling.	`STELLAR`

Recipe¶

import jax, jax.numpy as jnp
from jaxstro.units import STELLAR
from progenax import PlummerProfile, PlummerVelocityDF, build_spatial_ic

G = STELLAR.G
key = jax.random.PRNGKey(0)
masses = jnp.ones(2000)

# 1. Synthetic "observed" cluster at the true half-mass radius.
r_h_true = 1.3
obs = build_spatial_ic(PlummerProfile(r_h=r_h_true), masses,
                       PlummerVelocityDF(r_h=r_h_true), key=key, G=G, Q=0.5)
target = jnp.sqrt(jnp.mean(jnp.sum(obs.positions**2, axis=1)))   # RMS radius

# 2. Forward map: rebuild the IC at a trial r_h and measure the same observable.
def predicted_rms(r_h):
    ic = build_spatial_ic(PlummerProfile(r_h=r_h), masses,
                          PlummerVelocityDF(r_h=r_h), key=key, G=G, Q=0.5)
    return jnp.sqrt(jnp.mean(jnp.sum(ic.positions**2, axis=1)))

def loss(r_h):
    return (predicted_rms(r_h) - target) ** 2

grad_loss = jax.grad(loss)              # d(loss)/d(r_h) THROUGH the IC builder
print(f"loss(1.0) = {loss(1.0):.4e}, grad = {grad_loss(1.0):+.4f}")

# 3. Gradient descent through the generator (frozen key -> deterministic).
r_h = 1.0
for step in range(80):
    r_h = r_h - 0.1 * grad_loss(r_h)
print(f"recovered r_h = {r_h:.4f}  (truth {r_h_true})")

Verified output¶

Measured (PRNGKey(0), $N=2000$ , 80 GD steps, learning rate 0.1):

loss(1.0) = 4.9163e-01, grad = -3.2775
recovered r_h = 1.3000  (truth 1.3)

The gradient is non-zero and points toward the truth; gradient descent converges to the exact $\rh^{\rm true}$ .

Inputs and assumptions¶

Recipe¶

Verified output¶

See also¶