feat(parameterization): add constrained transforms for age and metallicity

Author: Tobias Buck

docs/development/gradient_production_pr_plan.md

@@ -15,14 +15,14 @@ proof-of-concept notebooks to production full-IFU workflows.
 
 ## PR stack
 
-1. `feat(inference-api)` (current)
+1. `feat(inference-api)` (completed)
 - Add `rubix.inference` module with:
   - parameter application without mutating baseline `RubixData`
   - `forward`, `loss`, and `value_and_grad` API
 - Add unit tests for copy semantics, deterministic loss calls, and analytical
   gradient checks.
 
-2. `feat(parameterization)`
+2. `feat(parameterization)` (current)
 - Add constrained parameter transforms (age/metallicity first) and tests.
 
 3. `refactor(gradient-modes)`

rubix/inference/__init__.py

@@ -1,5 +1,26 @@
 """Inference helpers for gradient-based modeling workflows."""
 
 from .api import apply_params, forward, loss, value_and_grad
+from .parameterization import (
+    IdentityTransform,
+    ParameterTransform,
+    SigmoidBounds,
+    SoftplusLowerBound,
+    apply_transforms,
+    build_age_metallicity_transforms,
+    inverse_transforms,
+)
 
-__all__ = ["apply_params", "forward", "loss", "value_and_grad"]
+__all__ = [
+    "IdentityTransform",
+    "ParameterTransform",
+    "SigmoidBounds",
+    "SoftplusLowerBound",
+    "apply_params",
+    "apply_transforms",
+    "build_age_metallicity_transforms",
+    "forward",
+    "inverse_transforms",
+    "loss",
+    "value_and_grad",
+]

rubix/inference/parameterization.py

@@ -0,0 +1,148 @@
+from dataclasses import dataclass
+from typing import Mapping
+
+import jax
+import jax.numpy as jnp
+from beartype.typing import Any
+
+ParamsTree = Mapping[str, Mapping[str, Any]]
+TransformTree = Mapping[str, Mapping[str, "ParameterTransform"]]
+
+
+class ParameterTransform:
+    """Base class for constrained parameter transforms."""
+
+    def forward(self, unconstrained: Any) -> Any:
+        """Map unconstrained values to constrained space."""
+        raise NotImplementedError
+
+    def inverse(self, constrained: Any) -> Any:
+        """Map constrained values to unconstrained space."""
+        raise NotImplementedError
+
+
+@dataclass(frozen=True)
+class IdentityTransform(ParameterTransform):
+    """No-op transform."""
+
+    def forward(self, unconstrained: Any) -> Any:
+        return unconstrained
+
+    def inverse(self, constrained: Any) -> Any:
+        return constrained
+
+
+@dataclass(frozen=True)
+class SoftplusLowerBound(ParameterTransform):
+    """Transform values to ``(lower, inf)`` via softplus."""
+
+    lower: float
+    eps: float = 1e-8
+
+    def forward(self, unconstrained: Any) -> Any:
+        return self.lower + jax.nn.softplus(unconstrained) + self.eps
+
+    def inverse(self, constrained: Any) -> Any:
+        shifted = constrained - self.lower - self.eps
+        safe = jnp.maximum(shifted, jnp.asarray(self.eps, dtype=shifted.dtype))
+        return jnp.log(jnp.expm1(safe))
+
+
+@dataclass(frozen=True)
+class SigmoidBounds(ParameterTransform):
+    """Transform values to ``(lower, upper)`` via sigmoid."""
+
+    lower: float
+    upper: float
+    eps: float = 1e-8
+
+    def __post_init__(self):
+        if self.upper <= self.lower:
+            raise ValueError("upper must be strictly larger than lower")
+
+    def forward(self, unconstrained: Any) -> Any:
+        width = self.upper - self.lower
+        return self.lower + width * jax.nn.sigmoid(unconstrained)
+
+    def inverse(self, constrained: Any) -> Any:
+        width = self.upper - self.lower
+        x = (constrained - self.lower) / width
+        x = jnp.clip(x, self.eps, 1.0 - self.eps)
+        return jnp.log(x) - jnp.log1p(-x)
+
+
+def apply_transforms(
+    params: ParamsTree,
+    transforms: TransformTree,
+    direction: str = "forward",
+) -> dict[str, dict[str, Any]]:
+    """Apply transform tree to a parameter tree.
+
+    Args:
+        params (ParamsTree): Nested parameter dictionary.
+        transforms (TransformTree): Nested transform dictionary with the same
+            keys as ``params`` for transformed leaves.
+        direction (str, optional): Either ``"forward"`` (unconstrained ->
+            constrained) or ``"inverse"`` (constrained -> unconstrained).
+            Defaults to ``"forward"``.
+
+    Raises:
+        ValueError: If ``direction`` is invalid.
+
+    Returns:
+        dict[str, dict[str, Any]]: Transformed parameter dictionary.
+    """
+    if direction not in {"forward", "inverse"}:
+        raise ValueError("direction must be one of {'forward', 'inverse'}")
+
+    transformed: dict[str, dict[str, Any]] = {}
+    for component, fields in params.items():
+        transformed[component] = {}
+        component_transforms = transforms.get(component, {})
+        for field, value in fields.items():
+            transform = component_transforms.get(field, IdentityTransform())
+            if direction == "forward":
+                transformed[component][field] = transform.forward(value)
+            else:
+                transformed[component][field] = transform.inverse(value)
+
+    return transformed
+
+
+def inverse_transforms(
+    params: ParamsTree,
+    transforms: TransformTree,
+) -> dict[str, dict[str, Any]]:
+    """Apply inverse transforms to a parameter tree."""
+    return apply_transforms(params=params, transforms=transforms, direction="inverse")
+
+
+def build_age_metallicity_transforms(
+    age_lower: float = 0.0,
+    age_upper: float = 20.0,
+    metallicity_lower: float = 0.0,
+    metallicity_upper: float = 0.05,
+) -> dict[str, dict[str, ParameterTransform]]:
+    """Build default transform tree for age and metallicity optimization.
+
+    Args:
+        age_lower (float, optional): Lower age bound in Gyr. Defaults to 0.0.
+        age_upper (float, optional): Upper age bound in Gyr. Defaults to 20.0.
+        metallicity_lower (float, optional): Lower metallicity bound.
+            Defaults to 0.0.
+        metallicity_upper (float, optional): Upper metallicity bound.
+            Defaults to 0.05.
+
+    Returns:
+        dict[str, dict[str, ParameterTransform]]: Transform tree for
+            ``stars.age`` and ``stars.metallicity``.
+    """
+    return {
+        "stars": {
+            "age": SigmoidBounds(lower=age_lower, upper=age_upper),
+            "metallicity": SigmoidBounds(
+                lower=metallicity_lower,
+                upper=metallicity_upper,
+            ),
+        }
+    }

tests/test_inference_parameterization.py

@@ -0,0 +1,89 @@
+import jax.numpy as jnp
+import pytest
+
+from rubix.inference import (
+    IdentityTransform,
+    SigmoidBounds,
+    SoftplusLowerBound,
+    apply_transforms,
+    build_age_metallicity_transforms,
+    inverse_transforms,
+)
+
+
+def test_identity_transform_roundtrip():
+    transform = IdentityTransform()
+    values = jnp.array([-2.0, 0.0, 3.0])
+
+    constrained = transform.forward(values)
+    recovered = transform.inverse(constrained)
+
+    assert jnp.allclose(constrained, values)
+    assert jnp.allclose(recovered, values)
+
+
+def test_sigmoid_bounds_forward_and_inverse():
+    transform = SigmoidBounds(lower=0.0, upper=20.0)
+    unconstrained = jnp.array([-3.0, 0.0, 3.0])
+
+    constrained = transform.forward(unconstrained)
+    recovered = transform.inverse(constrained)
+
+    assert jnp.all(constrained > 0.0)
+    assert jnp.all(constrained < 20.0)
+    assert jnp.allclose(recovered, unconstrained, atol=1e-5, rtol=1e-5)
+
+
+def test_softplus_lower_bound_forward_and_inverse():
+    transform = SoftplusLowerBound(lower=0.0)
+    unconstrained = jnp.array([-6.0, 0.0, 4.0])
+
+    constrained = transform.forward(unconstrained)
+    recovered = transform.inverse(constrained)
+
+    assert jnp.all(constrained > 0.0)
+    assert jnp.allclose(recovered, unconstrained, atol=1e-5, rtol=1e-5)
+
+
+def test_apply_transforms_tree_roundtrip():
+    params = {
+        "stars": {
+            "age": jnp.array([-1.0, 0.0, 1.0]),
+            "metallicity": jnp.array([-2.0, 0.5, 2.0]),
+            "mass": jnp.array([1.0, 2.0, 3.0]),
+        }
+    }
+    transforms = build_age_metallicity_transforms(
+        age_lower=0.0,
+        age_upper=20.0,
+        metallicity_lower=0.0,
+        metallicity_upper=0.05,
+    )
+
+    constrained = apply_transforms(params, transforms, direction="forward")
+    recovered = inverse_transforms(constrained, transforms)
+
+    assert jnp.all(constrained["stars"]["age"] > 0.0)
+    assert jnp.all(constrained["stars"]["age"] < 20.0)
+    assert jnp.all(constrained["stars"]["metallicity"] > 0.0)
+    assert jnp.all(constrained["stars"]["metallicity"] < 0.05)
+    assert jnp.allclose(recovered["stars"]["age"], params["stars"]["age"], atol=1e-5)
+    assert jnp.allclose(
+        recovered["stars"]["metallicity"],
+        params["stars"]["metallicity"],
+        atol=1e-5,
+    )
+    assert jnp.allclose(recovered["stars"]["mass"], params["stars"]["mass"])
+
+
+def test_apply_transforms_raises_on_bad_direction():
+    params = {"stars": {"age": jnp.array([0.0])}}
+    transforms = build_age_metallicity_transforms()
+
+    with pytest.raises(ValueError, match="direction must be one of"):
+        apply_transforms(params, transforms, direction="bad")
+
+
+def test_sigmoid_bounds_raises_for_invalid_bounds():
+    with pytest.raises(ValueError, match="upper must be strictly larger than lower"):
+        SigmoidBounds(lower=1.0, upper=1.0)