Training & Design Optimization
====================================

Both parameter training (fit distribution/model parameters to data) and
design optimization (choose a design variable, e.g. a position or an
abatement fraction, to minimize/maximize an objective) reduce to the same
operation: gradient descent through the compiled graph via `Optax
<https://optax.readthedocs.io/>`_. The only difference is what plays the
role of "the parameter being optimized" and what plays "the objective."

Three training entry points, from narrowest to widest scope:

.. list-table::
   :header-rows: 1

   * - Function
     - Scope
   * - :func:`~prophys.engine.fit_distribution`
     - One standalone distribution (univariate **or multivariate**), no
       model needed
   * - :func:`~prophys.engine.calibrate`
     - A compiled model, against observations of one attribute
   * - :func:`~prophys.engine.finetune`
     - The **entire model jointly**, against observations of several
       attributes at once

Calibration (MLE / MAP)
---------------------------

.. code-block:: python

   import prophys as prp

   compiled = model.compile()
   result = prp.calibrate(
       compiled,
       "acceptance_drop",       # attribute name
       observations,            # array of observed values
       n_steps=300,
       learning_rate=0.05,
   )
   print(result.params)         # fitted parameters, constrained space
   print(result.final_loss)

Internally, :func:`~prophys.engine.calibrate` builds the negative
log-likelihood (optionally plus a prior term for MAP) as a function of the
model's **unconstrained** parameters — the same space
:meth:`Param.forward <prophys.symbolic.Param.forward>` maps back from — so
Adam never has to worry about box constraints or positivity; those are
enforced structurally by the parameterization itself.

.. image:: /_static/generated/gas_dispersion_calibration.png
   :width: 60%
   :align: center

Freezing and selecting parameters
--------------------------------------

Both :func:`~prophys.engine.calibrate` and
:func:`~prophys.engine.finetune` accept ``trainable=[...]`` (whitelist) or
``frozen=[...]`` (blacklist) of parameter names. Frozen parameters keep
their current values but still participate in the forward computation —
the standard workflow for staged fitting (fit the physics first, then the
noise model, then release everything for a final joint pass).

.. code-block:: python

   prp.calibrate(compiled, "complaints", records, trainable=["dose_k", "dose_x0"])
   prp.finetune(compiled, all_observations, frozen=["emission_rate"])

Finetuning the entire model
---------------------------------

:func:`~prophys.engine.finetune` maximizes the **joint** likelihood over
several attributes at once. Because attributes typically share upstream
parameters (a common emission rate, a shared monitoring sigma, a copula
correlation), joint fitting pools every observation's information into
every shared parameter — with per-attribute ``weights`` if some data
sources are more trustworthy than others:

.. code-block:: python

   result = prp.finetune(
       compiled,
       {"exposure_res0": exposure_obs, "complaint_res0": complaint_records},
       weights={"complaint_res0": 2.0},
   )

Learning distributions directly
------------------------------------

:func:`~prophys.engine.fit_distribution` trains the ``Param`` leaves of a
standalone distribution — no model required, and multivariate data works
the same way as univariate (the leading axis indexes observations):

.. code-block:: python

   # Univariate: recover Weibull parameters from samples
   w = prp.Weibull(scale=prp.Param("scale", init=5.0, transform="softplus"),
                  concentration=prp.Param("k", init=1.0, transform="softplus"))
   prp.fit_distribution(w, wind_measurements)

   # Multivariate: learn a full covariance via its Cholesky factor
   mvn = prp.MultivariateGaussian(
       mean=prp.Param("mu", shape=(2,), init=jnp.zeros(2)),
       cholesky=prp.Param("L", shape=(2, 2), init=jnp.eye(2)),
   )
   prp.fit_distribution(mvn, paired_measurements)

   # Mixtures: learn component weights (softmax-normalized logits)
   mix = prp.Mixture([prp.Gaussian(0.0, 1.0), prp.Gaussian(10.0, 1.0)],
                    logits=prp.Param("logits", shape=(2,), init=jnp.zeros(2)))
   prp.fit_distribution(mix, samples)

The Cholesky and copula parameterizations are chosen so *any*
unconstrained gradient step yields a valid covariance/correlation — no
projections, no clipping, no failed updates.

Design optimization
------------------------

.. code-block:: python

   result = prp.optimize(
       objective,               # a scalar Expr to minimize
       wrt=[turbine_positions],
       constraints=[min_spacing_violation],  # each must evaluate <= 0
       n_steps=200,
       learning_rate=1.0,
   )
   print(result.params["pos"])

Constraints are enforced via a quadratic penalty; per-parameter bounds
don't need a constraint at all, since a bounded :class:`~prophys.symbolic.Param`
already can't leave its valid range under any gradient step.

If the objective references a :class:`~prophys.domain.RandomVariable`
(e.g. optimizing for a specific wind condition rather than marginalizing
over it), bind it via ``extra_env``:

.. code-block:: python

   result = prp.optimize(
       objective,
       wrt=[abatement_param],
       extra_env={"wind_speed": 6.0, "wind_direction": jnp.deg2rad(35.0)},
   )

Using a compiled model in external JAX solvers
---------------------------------------------------

A :class:`~prophys.engine.CompiledModel` is a bundle of pure JAX functions —
nothing about it is tied to prophys' own Optax loops. Any quantity it
produces (an expectation, a log-likelihood, a CVaR) can be closed over in an
ordinary Python function of a parameter pytree and handed to *any*
JAX-gradient-based solver; gradients flow through the native kernels'
``custom_vjp`` rules exactly as they do internally:

.. code-block:: python

   import jax

   def objective(raw_params):
       # any scalar the model can produce; raw_params is unconstrained space
       return -compiled.log_prob("exposure", observations, params=raw_params).sum()

   value_and_grad = jax.jit(jax.value_and_grad(objective))

   # jaxopt
   import jaxopt
   result = jaxopt.LBFGS(fun=objective).run(compiled.default_params())

   # optimistix
   import optimistix as optx
   sol = optx.minimise(lambda p, _: objective(p), optx.BFGS(rtol=1e-6, atol=1e-6),
                       compiled.default_params())

   # scipy, via JAX-computed gradients
   from scipy.optimize import minimize
   import numpy as np
   flat0 = np.asarray(compiled.default_params()["sigma"])
   # ... flatten/unflatten with jax.flatten_util.ravel_pytree as usual

The same applies to symbolic expressions directly: ``expr.evaluate(env)`` is
a pure function of ``env``, so ``jax.grad(lambda v: expr.evaluate({"x": v}))``
works with any downstream solver, inside or outside ``jit``.

Diagnostics
--------------

:func:`~prophys.engine.finite_difference_check` compares
``jax.grad`` against central finite differences — useful whenever a
:class:`~prophys.symbolic.Surrogate` wraps hand-written physics, to catch a
silent mismatch between the analytic and autodiff gradient.
:func:`~prophys.engine.check_identifiability` is a cheap heuristic flag for
"the loss was still improving meaningfully when optimization stopped."

.. autofunction:: prophys.engine.calibrate
.. autofunction:: prophys.engine.finetune
.. autofunction:: prophys.engine.fit_distribution
.. autofunction:: prophys.engine.optimize
.. autofunction:: prophys.engine.finite_difference_check
.. autofunction:: prophys.engine.check_identifiability
.. autoclass:: prophys.engine.CompiledModel
   :members: