Example: Gas Dispersion & Community Health Risk
===================================================

This chapter builds a complete prophys model step by step: an industrial
flare stack emits a pollutant, wind carries it towards nearby residences,
and we want (1) the exposure and complaint risk at each residence, (2) a
dose-response curve calibrated against observed complaints, and (3) the
emissions-abatement fraction that minimizes expected complaints while the
site stays compliant with its exclusion zone. The full runnable source is
``examples/gas_dispersion/model.py``.

.. image:: /_static/generated/gas_dispersion_risk_map.png
   :width: 80%
   :align: center

Step 1 — The frame and the geometry
--------------------------------------

Everything spatial lives in one shared coordinate system, the ``site``
frame. The scene is: a stack at the origin, four residences, a rectangular
regulatory exclusion zone, two neighbouring plant footprints, and a
pipeline corridor.

.. code-block:: python

   import jax.numpy as jnp
   import prophys as prp

   site = prp.Frame("site", units="m")

   stack = prp.Point(jnp.array([0.0, 0.0]), site)
   residences = prp.PointList(
       jnp.array([[300.0, 150.0], [500.0, -200.0], [800.0, 400.0], [150.0, -350.0]]), site
   )
   exclusion_zone = prp.Polygon(
       jnp.array([[-100.0, -100.0], [250.0, -100.0], [250.0, 250.0], [-100.0, 250.0]]), site
   )
   pipeline = prp.LineList(jnp.array([[[-1000.0, 100.0], [1000.0, 100.0]]]), site)

A 3D underground storage tank is modeled as a convex
:class:`~prophys.structures.Polyhedron` in halfspace form
:math:`\{x : Ax \le b\}` — a box centered at ``(100, 50, -5)`` becomes six
halfspaces (one per face):

.. code-block:: python

   tank_frame = prp.Frame("site_3d", ndim=3, units="m")
   center, half = jnp.array([100.0, 50.0, -5.0]), jnp.array([20.0, 15.0, 8.0])
   A = jnp.concatenate([jnp.eye(3), -jnp.eye(3)], axis=0)
   b = jnp.concatenate([center + half, -(center - half)])
   storage_tank = prp.Polyhedron(A, b, tank_frame)

Finally, terrain elevation is a raster :class:`~prophys.structures.Field`
(a gentle slope plus a ridge near the residences). Interpolating it is
differentiable, which matters in Step 3.

.. code-block:: python

   terrain = prp.Field(terrain_values, site, origin=(-1000.0, -1000.0), cell=(51.3, 51.3))

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

Step 2 — Initial uncertainty: wind as random inputs
------------------------------------------------------

The defining feature of this model is that uncertainty enters at the
*inputs*, not just as noise on the output. Wind speed follows a Weibull
distribution (the standard meteorological choice) and wind direction a
von Mises distribution (the circular analog of a Gaussian — a plain
Gaussian would be wrong because 0° and 360° are the same direction):

.. code-block:: python

   wind_speed = prp.RandomVariable("wind_speed", prp.Weibull(scale=6.0, concentration=2.0))
   wind_direction = prp.RandomVariable("wind_direction", prp.VonMises(loc=prp.deg2rad(35.0), kappa=3.0))

Every quantity computed downstream of these two leaves is itself a random
quantity. The compiled model marginalizes them out by Monte Carlo (see
:ref:`marginalization`), and — because both distributions sample by
reparameterization — gradients flow through the marginalization back to
any trainable parameter, including the wind distribution's own ``scale``
and ``kappa``.

Step 3 — The physics: a Gaussian plume, symbolically
--------------------------------------------------------

Ground-level concentration follows the standard Gaussian atmospheric
dispersion model. First, each receptor's position is rotated into
plume-aligned coordinates using the (random!) wind direction
:math:`\theta`:

.. math::

   x_d = \Delta x \cos\theta + \Delta y \sin\theta, \qquad
   x_c = -\Delta x \sin\theta + \Delta y \cos\theta

.. code-block:: python

   delta = receptors.coords - stack.coords
   dx, dy = delta[:, 0], delta[:, 1]
   downwind  =  dx * prp.cos(wind_direction) + dy * prp.sin(wind_direction)
   crosswind = -dx * prp.sin(wind_direction) + dy * prp.cos(wind_direction)

The plume's spread grows with downwind distance. Rather than hard-coding a
formula, the Pasquill–Gifford-style dispersion coefficients are supplied
as :class:`~prophys.transformations.TableLookup1D` characteristic curves —
interpolated, differentiable lookup tables:

.. code-block:: python

   sigma_y = prp.TableLookup1D([50, 200, 500, 1000, 2000], [8, 25, 55, 100, 180])(downwind_pos)
   sigma_z = prp.TableLookup1D([50, 200, 500, 1000, 2000], [5, 15, 30, 50, 80])(downwind_pos)

The effective release height is terrain-corrected by interpolating the
elevation field at the source and at each receptor — this is where
``Field.interp`` enters the physics:

.. code-block:: python

   eff_height = stack_height + (terrain.interp(stack) - terrain.interp(receptors))

Putting it together, the classic plume equation — with the emission rate
:math:`Q` reduced by a *trainable, bounded* abatement fraction (the design
variable of Step 6):

.. math::

   C = \frac{Q}{2\pi u\, \sigma_y \sigma_z}
     \exp\!\left(-\frac{x_c^2}{2\sigma_y^2}\right)
     \exp\!\left(-\frac{H_{\mathrm{eff}}^2}{2\sigma_z^2}\right)

.. code-block:: python

   abatement = prp.Param("abatement", init=0.0, bounds=(0.0, 0.9))
   Q = (1.0 - abatement) * prp.Param("emission_rate", init=5.0)

   conc = (Q / (2 * jnp.pi * u * sigma_y * sigma_z)) \
          * prp.exp(-0.5 * (crosswind / sigma_y) ** 2) \
          * prp.exp(-0.5 * (eff_height / sigma_z) ** 2)
   conc = prp.where(downwind > 0, conc, 1e-6)   # upwind receptors see no plume

Every line above is symbolic — ``jax.grad`` differentiates through the
rotation, the table lookups, the field interpolation, and the exponentials
in one pass.

Step 4 — Transformations: from concentration to risk
--------------------------------------------------------

Physics gives a concentration; the *questions* are about people and
regulators. Two transformations bridge that gap. A
:class:`~prophys.transformations.PiecewiseLinear` dose-response curve maps
concentration to complaint probability (``smooth`` keeps its gradient
continuous across the knots), and a
:class:`~prophys.transformations.Logistic` maps the exclusion-zone
compliance margin — a :meth:`Polygon.distance
<prophys.structures.Polygon.distance>` — to a regulatory-acceptance
probability:

.. code-block:: python

   dose_response = prp.PiecewiseLinear(
       knots=[0.0, 3e-5, 1e-4, 3e-4], values=[0.0, 0.05, 0.35, 0.85], smooth=1e-5
   )
   complaint_prob = prp.clip(dose_response(conc), 0.0, 1.0)

   compliance_margin = exclusion_zone.distance(residences)
   acceptance = prp.Logistic(x0=0.0, k=0.5, L=1.0)(compliance_margin)

One practical rule visible here: the ``smooth`` temperature must be small
relative to the knot spacing, and a probability fed into a
:class:`~prophys.distributions.Bernoulli` should be clipped to
:math:`[0, 1]` regardless.

Step 5 — Uncertain attributes and their correlation
-------------------------------------------------------

The named, observable quantities of the model are
:class:`~prophys.domain.UncertainAttribute` objects. Exposure at a
residence is LogNormal around the physical concentration (multiplicative
monitoring noise, with a *learnable* sigma); a complaint is a Bernoulli
draw from the dose-response output:

.. code-block:: python

   noise_sigma = prp.Param("noise_sigma", init=0.3, transform="softplus")

   exposure_0 = prp.UncertainAttribute(
       "exposure_res0",
       prp.LogNormal(mu=prp.log(prp.clip(conc_res0, 1e-6, None)), sigma=noise_sigma),
       unit="ug/m3",
   )
   complaint_res0 = prp.UncertainAttribute("complaint_res0", prp.Bernoulli(prob=complaint_prob_res0))

Residences 0 and 1 share the same weather, so their exposures are
correlated *even conditional on the mean model*. That is captured with a
:class:`~prophys.distributions.Correlation` (a bivariate Gaussian copula)
whose correlation parameter is itself trainable:

.. code-block:: python

   interactions = prp.AttributeInteractions(
       [exposure_0, exposure_1],
       prp.Correlation(exposure_0.distribution, exposure_1.distribution,
                      rho_raw=prp.Param("rho_raw", init=0.5)),
   )
   model = prp.ProbabilityModel(exposure_0, exposure_1, complaint_res0, interactions=interactions)

.. note::

   Attributes that get marginalized over Monte-Carlo wind samples are
   built from a **single-point** receptor (``conc_res0`` above), so the
   receptor batch dimension broadcasts cleanly against the
   ``(n_samples,)`` wind draws. The full four-residence concentration is
   computed separately, once, for the deterministic snapshot and the risk
   map.

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

Step 6 — Compile, train, optimize, export
---------------------------------------------

Compiling freezes the model into a stable interface
(``log_prob``/``sample``/``expectation``/``cvar``/…). Calibration fits the
dose-response against observed complaint records; here only the relevant
parameters are trained and the rest frozen:

.. code-block:: python

   compiled = model.compile(mode="opt", tau=5.0, n_samples=256)

   cal = prp.calibrate(compiled, "complaint_res0", complaint_records,
                      n_steps=200, learning_rate=0.05)

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

Design optimization then picks the abatement fraction. Note the division
of labor: the *box* constraint (abatement in :math:`[0, 0.9]`) is enforced
structurally by the bounded ``Param``, so only the cross-cutting
compliance constraint needs a penalty term. The wind condition is bound to
a design scenario via ``extra_env`` (design conditions on a scenario;
expectations marginalize):

.. code-block:: python

   result = prp.optimize(
       objective=total_complaint_prob,               # scalar Expr
       wrt=[abatement],
       constraints=[5.0 - compliance_margin],        # each must end up <= 0
       extra_env={"wind_speed": 6.0, "wind_direction": jnp.deg2rad(35.0)},
       n_steps=150, learning_rate=0.5,
   )

Finally, the calibrated model exports to the neutral
:class:`~prophys.export.ModelPackage` format for downstream consumers:

.. code-block:: python

   pkg = compiled.export()
   pkg.save("gas_model.json")

Visualizing the probability landscape
------------------------------------------

Because every quantity in the model is an ordinary symbolic expression,
evaluating it isn't limited to the four residences — evaluating it over a
dense :class:`~prophys.structures.Grid` turns any attribute into a
landscape plot. The complaint-probability landscape pushes the
concentration field through the exact same ``dose_response`` transform
used for calibration:

.. code-block:: python

   grid_points = plot_grid.positions()
   grid_conc = plume_concentration(grid_points)
   complaint_landscape = prp.clip(dose_response(grid_conc), 0.0, 1.0)

.. image:: /_static/generated/gas_dispersion_complaint_landscape.png
   :width: 75%
   :align: center

The regulatory-acceptance landscape only depends on geometry (distance to
the exclusion zone), so it is wind-independent and sharply bounded by the
zone itself — ``Logistic(x0=0)`` reads exactly 0.5 on the boundary:

.. code-block:: python

   compliance_landscape = exclusion_zone.distance(grid_points)
   acceptance_landscape = prp.Logistic(x0=0.0, k=0.5, L=1.0)(compliance_landscape)

.. image:: /_static/generated/gas_dispersion_acceptance_landscape.png
   :width: 75%
   :align: center

Primitives exercised
-----------------------

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

   * - Category
     - Used for
   * - :class:`~prophys.structures.Point` / :class:`~prophys.structures.PointList`
     - The stack (source) and four residences (receptors)
   * - :class:`~prophys.structures.Polygon` / :class:`~prophys.structures.PolygonList`
     - Regulatory exclusion zone; neighbouring plant footprints
   * - :class:`~prophys.structures.LineList`
     - Pipeline corridor distance
   * - :class:`~prophys.structures.Polyhedron`
     - 3D underground storage-tank exclusion volume
   * - :class:`~prophys.structures.Field` / :class:`~prophys.structures.Grid`
     - Terrain correction; dense risk-map evaluation
   * - :class:`~prophys.distributions.Weibull`, :class:`~prophys.distributions.VonMises`
     - Wind speed/direction as :class:`~prophys.domain.RandomVariable` inputs
   * - :class:`~prophys.transformations.TableLookup1D`, :class:`~prophys.transformations.PiecewiseLinear`, :class:`~prophys.transformations.Logistic`
     - Dispersion coefficients; dose-response; regulatory acceptance
   * - :class:`~prophys.distributions.LogNormal`, :class:`~prophys.distributions.Bernoulli`, :class:`~prophys.distributions.Correlation`
     - Exposure noise; complaint likelihood; shared-meteorology coupling
   * - :func:`~prophys.engine.calibrate`, :func:`~prophys.engine.optimize`, :class:`~prophys.export.ModelPackage`
     - Training, design, and export

Running it
-----------

.. code-block:: bash

   python examples/gas_dispersion/model.py