Unit 9: Project Specification
The capstone project, written out concretely. The science behind every term sits in Unit 8; the solution — finite-difference reference, forward PINN, inverse PINN — is in Unit 10. This unit is the spec: every equation, every parameter, every toy task you’ll need to reproduce the workshop.
9.1 The synthetic scenario
You are a fictional oceanographer with three thermistor chains deployed along a cross-shelf transect of the central Great Barrier Reef, offshore from Townsville. Each chain samples temperature at five depths, hourly, for 30 days. A storm passes through on day 10, and across the three sites the deeper sensors record signatures that do not line up with each other — one site cools quickly and recovers, another shows a deep slow cooling, and the third looks almost untouched at the surface.
| Site | Name | Lat, Lon | Depth | Sensor depths (m) | Dominant regime |
|---|---|---|---|---|---|
| A | Cleveland Bay | 19.20°S, 146.81°E | 15 m | 1, 4, 8, 12, 14 | tidal mixing, \mathrm{Pe}\ll 1 |
| B | Davies Reef | 18.83°S, 147.647°E | 60 m | 2, 10, 25, 45, 58 | classical thermocline, \mathrm{Pe}\sim 1 |
| C | Myrmidon Reef | 18.27°S, 147.39°E | 100 m | 2, 15, 40, 70, 95 | advection-dominated, \mathrm{Pe}\gg 1 |
All three sites are AIMS-monitored locations on the central GBR; the coordinates are sensible (Davies Reef’s real AIMS tower sits at 18°49'31''S, 147°38'50''E) but the time series are synthetic — generated by running a per-site finite-difference column model through a common storm scenario with \sigma = 0.05\,°\mathrm{C} observation noise. Source: units/unit_10/scripts/generate_mooring_csvs.jl; data: units/unit_10/data/mooring_{A,B,C}.csv plus units/unit_10/data/sites_metadata.csv.
Three hypotheses for the deeper-sensor cooling, applied at each site:
- The storm intensified upwelling, pumping cold water up from below.
- The storm intensified vertical mixing, drawing heat away from the surface faster.
- The storm caused reduced surface heating (cloud cover, evaporation).
You want a model simple enough to test these hypotheses against each mooring’s data. The 1D column below — coupled to a 2D shallow-water driver — is that model. The capstone solves it three times, once per site, with each site giving a qualitatively different fingerprint of the same event.
Tasks A and B don’t recover three separate drivers — they recover one scalar storm-forcing envelope \tau(t). Physically that envelope is the storm’s wind stress, which acts on the column several ways at once: it pumps water up (Ekman-driven upwelling w), stirs the surface layer (mixing \kappa_m), and dims the sunlight (cloud cover, Q_{\text{SW}}). The capstone’s toy column collapses those into a single forcing that enters as a heating source S(\zeta)\,\tau(t) — a deliberate simplification that keeps the inverse well-posed. So what the inverse recovers is the storm’s strength over time, not the separate upwelling / mixing / cloud responses.
So the recovered \hat\tau(t) doesn’t by itself answer “which hypothesis dominated”. The mechanism-discrimination step is the partition deliverable of Task B (§9.10): run the calibrated forward model with \hat\tau, then decompose the cooling at each site’s diagnostic depth (near its local thermocline) into the three terms
\underbrace{\Delta T_{\text{adv}}(t)}_{\int -w\,\partial_z T\, dt} \;+\; \underbrace{\Delta T_{\text{mix}}(t)}_{\int \partial_z(\kappa\,\partial_z T)\, dt} \;+\; \underbrace{\Delta T_{\text{flux}}(t)}_{\int \mathcal{S} + \text{surface BC}\, dt} \;=\; \Delta T_{\text{total}}(t),
and report which integrated contribution dominates over the storm window. The single-driver inversion plus the partition plot is what turns “we recovered \tau” into a statement like “the storm cooling here was mostly upwelling, with a smaller mixing contribution and little change in surface heating” — see Unit 10 §10.3 for the partitions the worked solution actually recovers. Task A stops at recovering \hat\tau (its timing and amplitude): a single mooring shows one column’s mechanism mix, but the §9.1 question is really about the contrast across regimes, which only Task B’s three-site study resolves.
If you want a direct mechanism-discrimination inverse instead (recover \tau_w, \tau_\kappa, \tau_Q separately), that’s a genuine extension and a stretch goal — see the “open questions” at the end of §9.10.
The simplest reading of the capstone is three independent inverse problems: each site recovers its own local wind-stress envelope \tau_i(t) from its own mooring data. That’s what §9.9 (Task A) below sets up.
An extended problem couples all three sites by sharing a single storm-event forcing \tau(t) (an atmospheric input felt across the whole ~100 km transect), and recovers it jointly from all three moorings’ data. The joint problem is better posed (three complementary regimes constrain one signal) and tells one coherent story instead of three parallel ones. Both versions are in scope for the workshop — the decoupled version is the baseline; the joint version is a stretch goal documented as a closing exercise.
9.2 The coupled model
The 1D column is the central object. A 2D shallow-water solver lives alongside it to drive the column with a realistic vertical velocity, but does not itself receive any feedback from the column (the SWE solve is offline). The three pieces:
1D vertical heat transport
Imagine a vertical pipe of seawater, 100 m tall, sitting under the sea surface. Heat enters from the top (sun, atmosphere). Cold water can be pushed up from the bottom (currents elsewhere). Inside the pipe, turbulence stirs warm and cold water together. We track only one thing: the temperature T as a function of depth z and time t. We ignore everything horizontal — no horizontal currents inside the pipe, no horizontal temperature variation. One column, by itself.
2D shallow water for horizontal flow
The linearised shallow water equations on a 2D domain produce horizontal velocity \mathbf{u}_h(x, y, t) in response to a wind-stress forcing. We use these as the driver of the column’s upwelling, not as a coupled-PDE PINN problem. The SWE solve is offline; its output is a time series at the mooring location.
Coupling: vertical velocity from horizontal divergence
First, a reminder of what we are after. w(z, t) is the vertical velocity of the seawater — how fast it moves up or down, in m/s — at depth z and time t. It is the one quantity through which the horizontal ocean reaches the 1-D column: it is the advection term (b), w\,\partial_z T, in the column equation of §9.3, and w > 0 (upwelling) drags cold deep water up past the sensors.
Here is the catch, and the thing worth seeing clearly: a 2-D shallow-water model never solves for w. Its unknowns are the horizontal current \mathbf{u}_h = (u, v), taken depth-uniform (barotropic). We recover w from mass conservation. Seawater is effectively incompressible, so the three velocity components cannot let fluid accumulate anywhere:
\partial_x u + \partial_y v + \partial_z w = 0 \qquad\Longleftrightarrow\qquad \partial_z w = -\,\nabla_h\cdot \mathbf{u}_h .
The physical reading is the whole idea in one line: where the horizontal flow converges (\nabla_h\cdot \mathbf{u}_h < 0), the water has nowhere to go but up — that vertical escape is the upwelling w.
To get the profile, integrate that relation upward from the seafloor, where the solid bottom forces w(-H) = 0. Because the SWE flow is depth-uniform, \nabla_h\cdot \mathbf{u}_h does not depend on z and the integral is immediate:
w(z, t) = -(z + H)\,\nabla_h\cdot \mathbf{u}_h(t).
So w is linear in z: zero at the bottom, growing in a straight line to its largest value at the surface, and rescaled at every instant by the single number \nabla_h\cdot \mathbf{u}_h(t) the 2-D model reports at the mooring. The full chain is therefore: wind stress → 2-D SWE → horizontal current \mathbf{u}_h(t) at the mooring → take its divergence → w(z, t) fed into the column. You build exactly this formula, by hand, in the exercise below.
The shipped capstone does not make you run a shallow-water solve. The column scripts take w as a prescribed storm envelope (toy scenario 4 of §9.11); replacing it with a genuine SWE-derived w(z, t) through the formula above is the scenario-5 design written up in Unit 10 §10.1. The committed units/unit_10/data/mooring_{A,B,C}.csv (generated once by units/unit_10/scripts/generate_mooring_csvs.jl, which calls the finite-difference column solver units/unit_10/scripts/column_fd.jl) are realistic records to look at (Exercise 9.1). The array each inverse task actually inverts is one it generates itself from a planted, known \tau^\star(t), so the recovery can be scored against the truth — the twin experiment of §9.7. No SWE run, and no committed CSV, enters the solution.
9.3 The column equation
\underbrace{\frac{\partial T}{\partial t}}_{\text{(a) change in time}} \;+\; \underbrace{w(z,t)\,\frac{\partial T}{\partial z}}_{\text{(b) advection}} \;=\; \underbrace{\frac{\partial}{\partial z}\!\left(\kappa(z,t)\,\frac{\partial T}{\partial z}\right)}_{\text{(c) diffusion}} \;+\; \underbrace{\mathcal{S}(z,t)}_{\text{(d) sun heating}}
on the domain z \in [-H, 0], t \in [0, T_f], with z = 0 at the surface and z = -H at the bottom. The four terms (a)–(d) were unpacked term-by-term in Unit 8 §8.3–§8.4; here we’ll just write them down.
w = \mathrm{d}z/\mathrm{d}t. With z measured upward from the surface (so z < 0 inside the ocean), w > 0 is upward motion of water. Upwelling — cold deep water rising — therefore corresponds to w > 0. The advection term w \, \partial_z T (on the left of the column equation above) then has the right physical sign; moved to the right it reads -w\,\partial_z T, the form Unit 8 §8.3 uses. The same note appears in Unit 8 — keep the two conventions aligned.
9.4 Boundary conditions
Top, z = 0: surface heat flux
\kappa\,\frac{\partial T}{\partial z}\bigg|_{z=0} \;=\; \frac{Q_{\text{np}}(t)}{\rho_0\, c_p}.
This is a Neumann (prescribed-flux) condition: it fixes the temperature gradient at the surface, not the temperature itself. The downward diffusive heat flux into the ocean from the air–sea interface equals the non-penetrating part of the surface heat budget, divided by the volumetric heat capacity of seawater (so units come out as K·m/s, matching \kappa\, \partial_z T). With z measured upward, \partial_z T > 0 corresponds to a warmer surface, and \kappa\, \partial_z T is the downward heat flux — positive Q_{\text{np}} heats the ocean by driving the surface temperature up.
- Q_{\text{np}}(t) — net non-penetrating heat flux at the surface. Sum of: longwave radiation, sensible heat, latent heat (evaporation). Units: W/m². Positive means heat into the ocean.
- \rho_0 — reference seawater density, \approx 1025\,\text{kg/m}^3.
- c_p — specific heat of seawater, \approx 3990\,\text{J/(kg}\cdot\text{K)}.
The penetrating shortwave Q_{\text{SW}} is not in this BC — it lives inside \mathcal{S}(z, t) as a body source instead.
Bottom, z = -H: cold reservoir
T(-H, t) \;=\; T_{\text{deep}}.
This is a Dirichlet (prescribed-value) condition — a fixed-temperature reservoir, as motivated in Unit 8 §8.5. The alternative, a zero-flux \partial_z T |_{-H} = 0 (itself a Neumann condition), gives an insulating floor instead.
Initial condition
T(z, 0) \;=\; T_0(z), \qquad T_0(z) = \tfrac12 (T_{\text{surface}} + T_{\text{deep}}) + \tfrac12 (T_{\text{surface}} - T_{\text{deep}})\, \tanh\!\bigl((z - z_t)/\delta_t\bigr),
The column starts out stratified, the way a real reef shelf does: a warm surface mixed layer at T_{\text{surface}} resting on cold deep water at T_{\text{deep}}, joined by a smooth transition — the thermocline (Unit 8 §8.5). The \tanh is simply the cleanest smooth step that does this. Reading the formula piece by piece:
- \tanh runs from -1 far below to +1 far above, so the profile swings from T_{\text{deep}} at depth (z \ll z_t, \tanh \to -1) up to T_{\text{surface}} near the surface (z \gg z_t, \tanh \to +1).
- The two prefactors turn that [-1, 1] swing into temperatures: \tfrac12(T_{\text{surface}} + T_{\text{deep}}) is the mean temperature the profile is centred on, and \tfrac12(T_{\text{surface}} - T_{\text{deep}}) is how far it reaches above and below that mean.
- z_t is the thermocline depth: at z = z_t the argument is zero and \tanh(0) = 0, so T_0(z_t) is exactly the midpoint temperature.
- \delta_t is the thermocline thickness: essentially all of the warm-to-cold change happens within a few \delta_t of z_t — a small \delta_t is a sharp thermocline, a large one a gradual one.
On the “\approx” question: z_t and \delta_t are values we set, not approximations of anything measured. For the synthetic scenario we fix z_t = -30\,\text{m} and \delta_t = 5\,\text{m} — a thermocline centred 30 m down and about 5 m thick, representative of the GBR shelf (all reference values are tabulated in §9.6; the per-site moorings vary this — Site B’s thermocline sits at -25 m).
9.5 Forcing functions
Surface heat flux
A minimal diurnal model:
Q_{\text{SW}}(t) \;=\; Q_{\text{SW}}^{\max}\, \max\!\bigl(0,\, \cos(2\pi t / \tau_d)\bigr), \qquad Q_{\text{np}}(t) \;=\; -\,Q_{\text{cool}}.
- \tau_d = 86400\,\text{s} — one day.
- Q_{\text{SW}}^{\max} — peak noon shortwave (e.g. 800\,\text{W/m}^2).
- Q_{\text{cool}} — steady net cooling from longwave + evaporation (e.g. 200\,\text{W/m}^2).
How to read the Q_{\text{SW}} term. The factor \cos(2\pi t/\tau_d) is the daily cycle: the 2\pi turns elapsed time into an angle, and dividing by \tau_d (one day) makes the period exactly one day, so the argument advances by 2\pi — one full revolution — every 24 h. The \max(0,\,\cdot) clips the negative half of the cosine to zero — the sun is simply off at night. Because the cosine peaks when its argument is 0, the model takes t = 0 to be local solar noon: Q_{\text{SW}} is at its peak Q_{\text{SW}}^{\max} there, falls to zero at t = \tau_d/4 (6 h later — “sunset”), stays zero through the night, and returns to its peak at t = \tau_d (next noon). It is therefore a deliberately crude 12 h-day / 12 h-night cycle, symmetric about noon, with no seasonal change in day length.
If your time axis is referenced to midnight instead (the usual convention for real forcing records), shift the phase: replace \cos(2\pi t/\tau_d) with \cos\bigl(2\pi(t - \tau_d/2)/\tau_d\bigr), which moves the noon peak to t = \tau_d/2 = 12 h and keeps it dark at t = 0. More generally \cos\bigl(2\pi(t - t_{\text{noon}})/\tau_d\bigr) puts solar noon at any chosen t_{\text{noon}}; the synthetic capstone data use the simplest choice, t_{\text{noon}} = 0.
Q_{\text{np}}, by contrast, carries no diurnal structure here — it is a constant net cooling -Q_{\text{cool}} (negative = heat leaving the ocean), a minimal stand-in for the longwave + evaporative loss that in reality also varies through the day.
Penetrating shortwave (the body source)
Sunlight is absorbed exponentially with depth (Beer–Lambert):
I(z, t) \;=\; Q_{\text{SW}}(t)\,e^{z/\zeta}, \qquad \mathcal{S}(z, t) \;=\; \frac{1}{\rho_0\, c_p}\,\frac{\partial I}{\partial z} \;=\; \frac{Q_{\text{SW}}(t)}{\rho_0\, c_p\, \zeta}\,e^{z/\zeta}.
The only z-dependence is the exponential, so differentiating it simply brings down a factor 1/\zeta. The heating \mathcal{S} is therefore largest right at the surface (z = 0) and decays with the same e-folding scale \zeta, depositing heat throughout the column rather than only at its top — which is exactly what makes penetrating sunlight a body source and not a surface flux. Integrated over a deep column it returns Q_{\text{SW}}(t)/(\rho_0 c_p): all of the penetrating sunlight ends up as heat in the water.
- \zeta — light penetration scale (e.g. \zeta \approx 10\,\text{m} for a single-band model; refined two-band Paulson–Simpson splits this into a 0.35 m red band and a 23 m blue–green band).
Vertical velocity (upwelling)
The column inherits w from the shallow-water driver, so its shape is fixed by the §9.2 derivation, not chosen freely. With a depth-uniform horizontal flow, mass conservation makes w linear in depth — pinned to zero at the seafloor and largest at the surface (§9.2):
w(z, t) \;=\; w_0(t)\,\frac{z + H}{H}, \qquad w_0(t) = -H\,\nabla_h\cdot\mathbf{u}_h(t),
so w_0(t) is the surface vertical velocity, set by the horizontal divergence the 2-D model reports at the mooring.
Why it is positive (upwelling). During the storm the surface current converges toward the mooring, \nabla_h\cdot\mathbf{u}_h(t) < 0; with the seafloor blocking any escape downward, the water has nowhere to go but up, so w_0(t) = -H\,\nabla_h\cdot\mathbf{u}_h(t) > 0 and the whole column lifts — drawing cold deep water past the sensors. A typical storm magnitude is w_0 \sim 10^{-5}\,\text{m/s}, about 1 m/day.
Eddy diffusivity
Three closures in increasing realism (pick one per experiment):
- Constant: \kappa = \kappa_0.
- Profile: \kappa(z) = \kappa_b + (\kappa_m - \kappa_b)\,e^{z/h_m} — large near the surface (mixed layer of scale h_m), small at depth.
- Stratification-dependent: \kappa(z, t) = \kappa_b + \kappa_0 / (1 + 5\,\mathrm{Ri})^2, with the Richardson number \mathrm{Ri} = N^2 / (\partial_z U)^2, buoyancy frequency N^2 = \alpha g\, \partial_z T, and a prescribed shear \partial_z U. Mixing is suppressed in stably stratified water (see Unit 8 §8.5).
9.6 Reference parameter values and dimensionless groups
| Symbol | Value | Units | Meaning |
|---|---|---|---|
| H | 100 | m | column depth |
| T_f | 2.592\times10^{6} (=30 d) | s | simulation horizon |
| \rho_0 | 1025 | kg/m³ | seawater density |
| c_p | 3990 | J/(kg·K) | specific heat |
| \alpha | 2\!\times\!10^{-4} | 1/K | thermal expansion |
| T_{\text{surface}} | 28 | °C | initial SST |
| T_{\text{deep}} | 18 | °C | deep reservoir |
| z_t,\, \delta_t | -30, 5 | m | thermocline depth, width |
| Q_{\text{SW}}^{\max} | 800 | W/m² | peak noon SW |
| Q_{\text{cool}} | 200 | W/m² | non-penetrating net cooling |
| \zeta | 10 | m | shortwave penetration scale |
| \kappa_b | 10^{-5} | m²/s | background diffusivity |
| \kappa_m | 10^{-3} | m²/s | mixed-layer diffusivity |
| h_m | 20 | m | mixed-layer scale |
| w_0 | 10^{-5} | m/s | peak upwelling (~1 m/day) |
The accompanying dimensionless groups — \mathrm{Pe} = w_0 H / \kappa_m \approx 1, T_\kappa = H^2/\kappa_m \approx 116 days, \tau_d / T_\kappa \sim 10^{-2} — are tabulated and interpreted in Unit 8 §8.6. (Naming caution: this Péclet number \mathrm{Pe} = w_0 H/\kappa_m appears in the non-dimensional PDE as the advection coefficient W_\text{adv}\approx1.2; the code’s variable Pe (=1) is the normalised diffusion-term coefficient, not the Péclet number.)
9.8 Which task is for which audience?
The capstone has one required task and one optional extension:
- Task A (§9.9) is the capstone everyone does. It runs end-to-end on a CPU and shows the full forward-PINN + inverse-PINN pipeline on a single site.
- Task B (§9.10) is an optional, advanced extension for those who want the operational-scale version — three sites jointly, the modern-PINN toolkit, and GPU scaling. Skip it entirely and you have still completed the capstone.
Task A runs end-to-end on a CPU in a few minutes and clears every §9.9.3 accuracy target there — you see the whole pipeline and learn everything it teaches without a GPU. The course GPU hub is available if you want to widen the network or iterate faster, but nothing in Task A needs it (see Unit 10 §10.2 for the measured CPU run). Task B, by contrast, really does want a GPU for its full three-site run — though even there the sub-scale prototypes run on CPU.
Pick by your hardware, time budget, and goals:
| Task A (§9.9) — required | Task B (§9.10) — optional / advanced | |
|---|---|---|
| Audience | every participant; self-paced learner, one-day workshop | research project, advanced student, industrial prototyping |
| Hardware | CPU (a few minutes); GPU optional | CPU for the sub-scale prototypes; NVIDIA GPU for the full three-site run |
| Time budget | ~4 min on a CPU | ~30 min on GPU per run; hours on CPU |
| Geometry | one site (Davies Reef, H = 60 m) | three sites jointly (Cleveland Bay + Davies Reef + Myrmidon Reef, H up to 100 m) |
| PINN toolkit | small MLP, static hand-tuned weights | + Fourier features, adaptive weights, causal training |
| Parameter count | ~5 000 | ~53 000 |
| What it demonstrates | the full forward-PINN + inverse-PINN pipeline | operational engineering — joint inversion, modern fixes, GPU scaling |
Both tasks share §9.7 (forward solve) and §9.11 toy ladder. The remainder of this unit is each task’s full spec.
9.9 Task A — single-site inverse (CPU; GPU optional)
The required capstone. Recover the storm wind-stress envelope \tau(t) from a single mooring’s data using a small PINN that trains in a few minutes on a CPU.
The goal is to see the full forward-PINN + inverse-PINN pipeline end-to-end. Everything is kept small enough that the training run finishes in coffee-break time on a CPU; the same code runs unchanged (and faster) on the course GPU hub if you have access — a GPU is optional here, not required.
9.9.1 Spec
- Site. Site B (Davies Reef) — H = 60 m, \mathrm{Pe} \sim 1: a 60 m column with vertical advection (upwelling) on top of diffusion, the classical thermocline of the three §9.1 sites, where the storm leaves a sharp, depth-structured fingerprint. 5 sensors at z = -2, -10, -25, -45, -58 m.
- Data. Synthetic Site-B mooring you generate yourself: a known \tau^\star(t) forward-solved, then sampled at the 5 sensor depths × 720 hourly steps with \sigma = 0.015\,°\text{C} Gaussian noise (a quality mooring; the committed CSVs use the noisier 0.05 °C) = 3 600 points — so the recovered \hat\tau scores against a known truth (twin experiment, §9.7). The committed
data/mooring_B.csvis the realistic look-only counterpart (Exercise 9.1). - Unknown. A single scalar function \tau(t) — the local wind-stress envelope — over the 30-day window.
- Given. Constant diffusivity \kappa_m (→ the diffusion number Pe), the upwelling profile w(\zeta)=\zeta (→ W_\text{adv}), the Beer–Lambert heating shape S(\zeta), the deep reservoir T_{\text{deep}}, and the reference parameters from §9.6 — which step 2 folds into Pe, W_\text{adv}, S(\zeta), the only numbers the solver uses.
- Networks. T_\theta(z, t) as a 4-layer 32-neuron MLP (\tanh); \tau_\phi(t) as a 32-neuron MLP. ~5 000 parameters total.
- Collocation. N_r = 4000 residual points, N_b = 800 BC points, N_d = 5 \text{ sensors} \times 720 \text{ samples} = 3600 data points.
9.9.2 Workflow
Follow these seven steps in order. Each maps directly onto a block of the starter skeleton below.
- Look at the data first. Read
data/mooring_B.csv(5 sensors × 720 hourly samples) and plot all five depths (§9.1 exercise) to see the storm’s fingerprint at this site. Treat this as orientation: the array you actually invert is one you generate yourself in step 4 from a known storm, so the recovery can be scored against the truth (a twin experiment, §9.7). (The CSV was produced once byscripts/generate_mooring_csvs.jl, which writes all three sites’ files.) - Build the column constants and forcings from §9.5–§9.6: the diffusivity \kappa(z), the upwelling profile w(z), the Beer–Lambert source shape, the initial profile T_0(z), and the deep reservoir T_{\text{deep}} — then non-dimensionalise into the handful of numbers the solver actually uses: \mathrm{Pe}, W_\text{adv}, the heating shape S(\zeta), and w(\zeta)=\zeta (see Unit 10 §10.2).
- Define the two networks. T_\theta(z, t) recovers the temperature field; \tau_\phi(t) recovers the unknown storm envelope. Bake the hard deep BC (T = T_{\text{deep}} at z = -H) and the hard IC into the field ansatz so those two constraints hold for free and only the surface flux BC stays a soft penalty (Unit 7 §7.3).
- Make your data with a known storm — and trust the forward solve first. Pick a known storm \tau^\star and run a forward solver to produce T(z,t); sampling its five sensor depths and adding noise gives the synthetic mooring you will invert. Use the
MethodOfLines.jlreference of §9.7, or — as the shipped solution does — a short hand-coded explicit finite-difference integrator for the same column. Because you planted \tau^\star, the recovered \hat\tau can be scored against a known truth — the twin experiment (§9.7). Both networks then train jointly (steps 5–6), not in two stages; the forward field is certified afterwards by checking T_\theta lands within the noise floor (target L^2 < 0.05\,°\text{C} against the reference). - Assemble the joint inverse loss \mathcal{L} = \lambda_r \mathcal{L}_{\text{PDE}} + \lambda_d \mathcal{L}_{\text{data}} + \lambda_b \mathcal{L}_{\text{surf-BC}} + \lambda_{\text{reg}} \int |\tau_\phi'|^2\,dt (the hard IC + deep BC mean no separate \mathcal{L}_{\text{IC}} / deep-BC terms). The loss weights are the whole game here — a large data weight \lambda_d and a tiny smoothing weight \lambda_{\text{reg}} are what keep the recovered storm peak from collapsing; see the weight-sweep table in Unit 10 §10.2.
- Train the inverse PINN. A few thousand steps already show the method working; ~12 000 Adam steps reach the §9.9.3 target (about four minutes on a CPU).
- Diagnose and report — overlay \hat\tau(t) on the truth and write the honest assessment of where the recovery fails. Optionally add the residual histogram vs t (Unit 7 §7.5) and a heat-budget closure check at each sensor depth.
Starter skeleton (Task A)
A scaffold for steps 3–6, distilled from the shipped solution task_a_inverse_pinn.jl. Fill in the # TODO (forward-solve a known storm to make your data, then non-dimensionalise — the twin experiment of §9.7, not a load of the committed CSV) from the equations in §9.3–§9.6; the full worked version — FD reference, metrics, and the validated weights — is in Unit 10 §10.2 behind the solution gate. It is for reading, not running.
§9.3–§9.6 set up the column in dimensional variables (z \in [-H,0] in metres, T in °C, a tanh stratified start, a surface cooling flux, the diurnal Q_\text{SW}). The skeleton below solves that same column after the change of variables step 2 of the workflow asks you to build (worked in full in Unit 10 §10.2):
\zeta = \frac{z + H}{H}, \qquad \tau_t = \frac{t}{T_f}, \qquad \tilde T = \frac{T - T_\text{deep}}{\Delta T},
with \zeta = 0 the deep bottom and \zeta = 1 the surface (the opposite orientation to z). This \zeta is the non-dimensional depth coordinate; it is not the §9.5 light-penetration scale, a different quantity that happens to reuse the symbol. Here \Delta T is the fixed surface-to-deep temperature scale (\approx 3\,°C at Site B — distinct from a timestep \Delta t; the §9.6 reference column has a larger \approx 10\,°C span). This gives the non-dimensional column the skeleton integrates,
\partial_{\tau} \tilde T = -W_\text{adv}\,w(\zeta)\,\partial_\zeta \tilde T + \mathrm{Pe}\,\partial_{\zeta\zeta}\tilde T + S(\zeta)\,\tau(t), \qquad \mathrm{Pe} = 1,\ \ W_\text{adv} = 1.2,\ \ w(\zeta) = \zeta,
where \partial_\tau is the derivative in non-dimensional time \tau_t = t/T_f and the standalone \tau(t) is the storm forcing — the two uses of \tau kept distinct as in §10.2.
Task A runs the clean source-only reduction: a flat start at T_\text{deep} and an insulating surface, so the §9.4 tanh initial profile, the surface cooling flux, and the diurnal cycle all fold out, and the storm enters only through S(\zeta)\,\tau(t) (the source the §9.1 callout already asserts). Because \tilde T is the anomaly from that flat reference, the hard IC \tilde T(\zeta,0) = 0 and deep BC \tilde T(0,\tau) = 0 baked into the ansatz \tilde T = \zeta\,\tau\,N are zero by construction — not a clash with the §9.4 stratified profile — and the surface condition is the insulating \partial_\zeta \tilde T(1,\tau) = 0.
using Lux, Optimisers, Zygote, Random, Statistics
# ── physics (non-dimensional column; see §9.3–§9.6 and Unit 10 §10.2) ──
const Pe = 1.0f0 # diffusion number for Site B
const Wadv = 1.2f0 # advection number (upwelling) → Péclet ~ 1
S_shape(ζ) = exp.((ζ .- 1f0) ./ 0.25f0) # known Beer–Lambert heating shape S(ζ)
w_shape(ζ) = ζ # upwelling: 0 at depth, max at surface
# TODO: forward-solve a known τ★ (FD reference, §9.7) → sample 5 depths + noise
# → (ζ_obs, τ_obs, T̃_obs) [non-dimensionalise z,t,T]. Twin experiment, not a CSV load.
# ── two networks ──────────────────────────────────────────────────────
T_net = Chain(Dense(2=>32, tanh), Dense(32=>32, tanh), Dense(32=>32, tanh), Dense(32=>1)) # field T̃_θ(ζ,t)
τ_net = Chain(Dense(1=>32, tanh), Dense(32=>32, tanh), Dense(32=>1)) # storm τ_φ(t)
# Hard IC (t=0 ⇒ 0) and hard deep BC (ζ=0 ⇒ 0) baked into the ansatz
# (T̃ = non-dimensional temperature; the NN weights are p, not θ):
T̃(p, ζ, t) = ζ .* t .* first(T_net(vcat(ζ, t), p, stT)) # T̃(0,·)=0, T̃(·,0)=0 for free
τφ(q, t) = first(τ_net(t, q, stτ)) # recovered forcing
# Derivatives by finite-difference-in-input (no nested AD; Unit 5 §5.3):
h = 2f-3
∂t(p,ζ,t) = (T̃(p,ζ,t.+h) .- T̃(p,ζ,t.-h)) ./ (2h)
∂ζ(p,ζ,t) = (T̃(p,ζ.+h,t) .- T̃(p,ζ.-h,t)) ./ (2h)
∂ζζ(p,ζ,t) = (T̃(p,ζ.+h,t) .- 2 .* T̃(p,ζ,t) .+ T̃(p,ζ.-h,t)) ./ h^2
# ── loss: the WEIGHTS are the lesson (large λ_d, tiny λ_reg) ───────────
λ_d, λ_b, λ_reg = 6000f0, 10f0, 1f-5
function loss(pT, pτ)
L_pde = mean(abs2, ∂t(pT,ζc,tc) .+ Wadv .* w_shape(ζc) .* ∂ζ(pT,ζc,tc)
.- Pe .* ∂ζζ(pT,ζc,tc) .- S_shape(ζc) .* τφ(pτ,tc))
L_data = mean(abs2, T̃(pT, ζ_obs, τ_obs) .- T̃_obs)
L_bc = mean(abs2, ( T̃(pT,ζ1,tb) .- T̃(pT,ζ1 .- h,tb) ) ./ h) # surface ∂ζT̃(1,·)=0
L_reg = mean(abs2, (τφ(pτ,tg.+h) .- τφ(pτ,tg.-h)) ./ (2h)) # H¹ smoothness on τ
return L_pde + λ_d*L_data + λ_b*L_bc + λ_reg*L_reg
end
# ── train (Adam on both networks) ─────────────────────────────────────
optT = Optimisers.setup(Adam(1f-3), pT)
optτ = Optimisers.setup(Adam(3f-3), pτ)
for _ in 1:N_iters
gT, gτ = Zygote.gradient(loss, pT, pτ)
optT, pT = Optimisers.update(optT, pT, gT)
optτ, pτ = Optimisers.update(optτ, pτ, gτ)
end- Forward before inverse. Get the field network matching the FD reference with a known \tau first. An inverse stacked on a broken forward is undebuggable.
- If the recovered peak is too low, raise \lambda_d before anything else (not the collocation count) — that is the single most important knob (Unit 10 §10.2).
- To run on the GPU, wrap parameters, inputs and data in
Lux.gpu_device()— the loss and training loop above are unchanged.
9.9.3 Success criteria
| Metric | Target |
|---|---|
| Forward PINN L^2 vs FD reference | < 0.05\,°\text{C} |
| Recovered \hat\tau peak-amplitude error | < 15% |
| Storm-day timing error | < 2 h |
9.9.4 Expected runtime
| Stage | Wall-clock (CPU) |
|---|---|
| Finite-difference reference solve (1 site, 30 days) | ~1 s |
| Forward-solve certification (optional, separate; §10.2) | ~2 min |
| Inverse PINN training (Adam, 12 000 steps) | ~4 min |
Total: under 10 minutes from a cold cache (measured on a CPU — the inverse run lands the recovered storm peak at ~10% error). Re-running to retune \lambda_d / \lambda_\text{reg} adds a few minutes per pass.
9.9.5 Deliverables
Required:
- Plot of \hat\tau(t) overlaid on the synthetic truth.
- A one-paragraph honest assessment of where the recovery fails (typically the rising edge — smoothed by the H^1 penalty — and the tail — over-relaxed).
Optional extras (good practice, not needed to clear the task):
- Residual-histogram-vs-time diagnostic (Unit 7 §7.5).
- Heat-budget closure plot for each sensor depth.
The mechanism partition — decomposing the cooling into advection, mixing, and source contributions to answer the §9.1 “which hypothesis dominated” question — is deferred to Task B (§9.10). A single site gives you only one column’s breakdown; the story the §9.1 question really asks for is the contrast — how the advection share climbs across the three regimes — which only the three-site study resolves (see Unit 10 §10.3).
9.9.6 What you don’t do here
- No cross-site coupling.
- No Fourier-feature embedding.
- No adaptive loss weighting.
Those land in the optional Task B. (A GPU, on the other hand, is welcome here — it just isn’t required; see the note in §9.8.)
9.10 Task B (optional, advanced) — three-site joint inverse
Task B is an advanced extension for those who want the operational-scale problem. You do not need it to complete the capstone — Task A (§9.9) is the required piece. Take Task B on if you have GPU access (the course hub) and want to see joint multi-site inversion at full scale (the modern-PINN toolkit is an optional scale-up, not needed to clear the bar). The sub-scale prototype steps below still run on a CPU.
The advanced capstone. Recover a single shared wind-stress envelope \tau(t) from all three moorings jointly, on the full H = 100 m / 30-day domain. CPU runtime is hours; GPU runtime is minutes. We develop the CPU sub-scale prototypes here and write the GPU-launch checklist.
Task B exploits the three mooring regimes — Cleveland Bay’s diffusion-only column pins the storm amplitude through its mixing response, Myrmidon’s advection-dominated column reveals the shape of the deep upwelling pulse, and Davies Reef sits between, anchoring the timing. A single coherent storm signal must explain all three at once.
9.10.1 Spec
- Sites. A + B + C jointly. Three column PINNs T^{(i)}_\theta(z, t), i \in \{A, B, C\}. Shared parameters for the wind-stress envelope \tau_\phi(t); separate parameters for each site’s temperature network.
- Data. Self-generated per site (twin experiment, §9.7): one shared known \tau^\star(t), three FD reference solves, each sampled at 5 sensors × 720 steps + noise = 3 × 3 600 = 10 800 points. The committed
data/mooring_{A,B,C}.csvare the look-only counterparts (Exercise 9.1). - Unknown. \tau(t) — one shared scalar function, three moorings’ worth of data constraining it.
- Networks. T^{(i)}_\theta(z, t) as a 6-layer 64-neuron MLP per site (~17 k parameters each); \tau_\phi(t) as a 3-layer 48-neuron MLP. ~53 000 parameters total.
- Modern PINN toolkit (Unit 7 §7.3) — optional scale-up. Fourier-feature embedding, gradient-balanced adaptive loss weighting, and causal time training are the menu for pushing to the deep H = 100 m column. The shipped 3-site run (
task_b_joint_inverse.jl) does not use them — it clears the bar with the static Task-A recipe (see §10.3 Step 6 and the tips below). The hard BC at z = -H and at the IC are always on. - Collocation. The shipped run uses N_r = 4\,000 per site; scaling to the deep column wants up to N_r = 50\,000. Data: N_d = 3\,600 per site (5 sensors × 720 hourly samples), 10 800 across the three sites.
9.10.2 Workflow
- Look at the three sites’ data first. Read the committed
data/mooring_{A,B,C}.csv(produced once byscripts/generate_mooring_csvs.jl) to see the three regimes (Exercise 9.1). As in Task A, treat this as orientation: the arrays you actually invert are ones you generate yourself in the prototype below, from a single known shared \tau^\star, so the joint recovery scores against a known truth (twin experiment, §9.7). - Sub-scale prototype on a single column. Run the static Task-A recipe (hard BC, static weights) on one site alone with N_r \approx 4\,000. A few minutes on CPU. Sanity check: it should match Task A’s single-site accuracy — the joint machinery is what earns its keep here, not any per-site fix.
- Two-site joint inverse on A + B (
task_b_subscale_prototype.jl). 32 × 4 networks, N_r = 4\,000 per site, static weights. ~7 min on CPU. Lets you see the joint-vs-decoupled improvement and tune cross-site loss weights. - GPU-launch checklist. The shipped 3-site run is essentially a device swap from the CPU prototype:
- wrap
ps, inputs, and data inLux.gpu_device(), - keep the same Adam-only schedule and static weights,
- (optional scale-up only) Fourier features, adaptive weights, or an L-BFGS polish if you deepen all three columns to H = 100 m,
- JAX equivalents (
jinns/Equinox) if your team prefers the JAX stack (Unit 7 §7.6).
- wrap
- Queue the full run on a GPU when available — the code doesn’t change, only the
deviceand the collocation count.
Starter skeleton (Task B)
Task B reuses Task A’s column, loss, and training loop almost verbatim — the one structural change is that all sites share a single \tau-network while each keeps its own temperature network. The scaffold below shows that extension; the worked two-site prototype is task_b_subscale_prototype.jl and the full three-site trainer is task_b_joint_inverse.jl (both in Unit 10 §10.3, behind the solution gate). Start from the Task A skeleton (§9.9.2) and add the shared-\tau loop. Display-only; not executed by render.
# Per-site column model: Site A is diffusion-only (Wadv=0); Site B and C
# add surface-intensified upwelling (Wadv>0) so their Péclet is larger.
struct Site; Pe::Float32; Wadv::Float32; ℓS::Float32; end
S_shape(s, ζ) = exp.((ζ .- 1f0) ./ s.ℓS)
w_shape(ζ) = ζ # linear surface-intensified upwelling, the §9.2 form w∝(z+H):
# 0 at depth (ζ=0), max at surface (ζ=1): linear w from the SWE (§9.2/§9.5; the toy scenario-3 figure uses sin)
sites = [Site(1.5f0, 0f0, 0.20f0), # A — Cleveland Bay (diffusion only)
Site(1.0f0, 1.2f0, 0.25f0)] # B — Davies Reef (advection + diffusion)
# ONE shared storm network; one temperature network PER site.
τ_net = Chain(Dense(1=>64, tanh), Dense(64=>64, tanh), Dense(64=>1)) # shared τ_φ(t)
T_nets = [Chain(Dense(2=>64, tanh), Dense(64=>64, tanh), Dense(64=>1)) for _ in sites]
# Per-site residual now carries the advection term too (cf. §9.3 term (b)):
function site_residual(s, T̃, ∂t, ∂ζ, ∂ζζ, τφ, ζc, tc)
∂t(ζc,tc) .+ s.Wadv .* w_shape(ζc) .* ∂ζ(ζc,tc) .- # advection
s.Pe .* ∂ζζ(ζc,tc) .- S_shape(s,ζc) .* τφ(tc) # diffusion + shared source
end
# Joint loss = Σ_sites (PDE + data + BC), + ONE H¹ prior on the shared τ:
function joint_loss(pτ, pTs)
L = 0f0
for (s, pT) in zip(sites, pTs)
L += site_pde(s, pT, pτ) + λ_d*site_data(s, pT) + λ_b*site_bc(s, pT)
end
return L + λ_reg * tau_smoothness(pτ)
end
# Train pτ together with every pT (one Adam state per parameter set), then
# compare the joint τ̂ against each site-alone inverse — the joint should win.- Same recipe, more sites. Keep Task A’s static weights (\lambda_d = 6000, \lambda_{\text{reg}} = 10^{-5}, hard IC + deep BC). The “modern toolkit” (Fourier features, adaptive weighting, causal training) is the menu for the deep H=100 m column — reach for it only once the static baseline is trustworthy (Unit 10 §10.3 explains why).
- Validate the headline claim. Run each site alone, then jointly. The joint \hat\tau peak error should beat both single-site recoveries — that cross-site gain is the point of Task B.
- GPU launch is a device swap, not a rewrite:
Lux.gpu_device()on parameters/data, optionallyReactant.@compilethe step, same Adam schedule (see step 4 above).
9.10.3 Success criteria
| Metric | Sub-scale on CPU | Full scale on GPU |
|---|---|---|
| Sub-scale prototype matches Task A | yes, < 1.2× error | — |
| Two-site joint \hat\tau peak error | ~7% | — |
| Full 3-site joint \hat\tau peak error | — | ~5–7% |
| Residual histogram monotonically decreasing in t, all three sites | — | yes |
9.10.4 Expected runtime
The course GPU hub is an NVIDIA A10G (g5.2xlarge); the worked solution’s measured times are in Unit 10 §10.3.
| Stage | CPU | GPU (A10G) |
|---|---|---|
| Single-column sub-scale prototype | ~5 min | ~1 min |
Two-site joint inverse (A + B, task_b_subscale_prototype.jl) |
~7 min | ~2 min |
Three-site joint inverse (task_b_joint_inverse.jl) |
~2 h | ~5 min |
| Optional production scale-up (128×6, N_r 50k) | — | ~1 h |
The CPU full-scale column is not a typo: technically possible but the iteration loop kills the development cycle. Prototype on CPU at reduced scale, then deploy on GPU at full scale.
9.10.5 Deliverables
- Sub-scale prototype results on Task A’s column with Task B’s architecture.
- Two-site joint-inverse \hat\tau(t) overlaid on the synthetic truth, with the per-site decoupled inverses on the same axes for comparison.
- Mechanism partition at all three sites (table or plots; z = -10\,\text{m} for Site A, z = -30\,\text{m} for Site B, z = -50\,\text{m} for Site C — each one near or just below the local thermocline). The partition should show qualitatively different mechanism weights across the three sites — that’s the §9.1 multi-site story made quantitative.
- The GPU-launch checklist as a
.mddocument. - A written assessment of what changes at the full 3-site, H = 100 m scale: predicted accuracy, predicted training cost on GPU, what could still go wrong (causality violation at Myrmidon, BC imbalance from the four-orders-of-magnitude scale gap between \kappa at surface and depth).
9.10.6 Open questions for the full GPU run
These are honest “we don’t know yet” items the participant should report on if they get GPU time:
- Will causality violation re-appear at Myrmidon Reef (advection-dominated, \mathrm{Pe} \gg 1)? Possibly need a per-site causal scheduler.
- Will gradient-balancing converge on a stable weight ratio? The static recipe is stable at sub-scale; gradient-balancing itself is untested here (§10.3 notes a naive adaptive scheme can run away).
- How well does the recovered \tau(t) correlate with the independent SWE-driver w(t) inferred from local wind observations? This is the validation step that turns the recovered synthetic answer into something an oceanographer trusts.
- Stretch goal — direct mechanism inversion. Instead of one shared \tau(t) implicitly coupling all three mechanisms, expand the inverse to recover three separate scalar functions \tau_w(t), \tau_\kappa(t), \tau_Q(t) (driving upwelling, mixing, and surface-flux modulation independently). This makes the §9.1 hypothesis discrimination direct rather than partition-inferred, at the cost of a three-fold-larger inverse problem (worse conditioning, more regularisation tuning). A good test of whether the joint three-site data is informative enough to separate three coupled drivers.
9.11 Toy-task ladder (shared)
Build up to the forward problem in five scenarios, each isolating one mechanism, so when the full storm lands you can read each contribution off the trace. These are exactly the five scenarios worked and plotted in Unit 10 §10.1 — the numbering matches one-to-one. Task A and Task B both start from these.
Pure diffusion, steady forcing. Set w = 0, \mathcal{S} = 0, \kappa constant, Q_{\text{np}} constant. Solve to steady state. Expected: a linear T(z) profile balancing surface flux against the deep reservoir. Sanity check: analytic solution exists.
Add the diurnal cycle. Turn on time-varying Q_{\text{SW}}(t) and the body source \mathcal{S}. Keep w = 0, \kappa constant. Expected: a diurnal warm layer in the top few metres that warms in the afternoon and erodes overnight. Bulk profile barely changes over 30 days.
Add upwelling. Turn on w(z) with w_0 = 10^{-5}\,\text{m/s}. Expected: cold water from depth invades the bulk; SST drops slowly over weeks. The diurnal warm layer survives but sits on top of cooler water. (The shipped toy scenarios keep \kappa constant at \kappa_m; the §9.5 depth-varying profile closure is taught but not wired into these illustrations.)
Storm fingerprint (prescribed gust). A Gaussian gust on day 10 (~3 days wide) modulating the upwelling w (×5 at peak) and Q_{\text{SW}} (cloud cover, ×0.5 at peak) with prescribed envelopes. Expected: deeper sensors cool first (upwelling fingerprint), surface stays close to baseline as cloud and stress partly cancel. This is the storm the capstone inverts.
The synthetic scenario from the SWE driver. Replace the prescribed envelopes of scenario 4 with a genuine 2-D shallow-water solve: a Gaussian wind gust over the mooring produces \mathbf{u}_h(t), hence w(t) (via §9.2) and Q_{\text{SW}}(t), fed into the column. This is the physically self-consistent storm — the headline extension, designed in Unit 10 §10.1, not shipped with the solution.