# Advection schemes

TRACMASS-JAX has the two trajectory schemes of Fortran TRACMASS. You choose
one with `scheme=` in `steady_state` and `time_dependent`:

| `scheme` | Fortran module | JAX factory |
|---|---|---|
| `"tstep"` (default) | `mod_pos_tstep.F90` | `piecewise(grid)` |
| `"analytical"` | `mod_pos_tanalytical.F90` | `analytical(grid, tseas)` |

The factories live in `tracmass.advection_scheme`. In the time-stepping
scheme the fluxes are frozen during a sub-step and interpolated in time
between sub-steps. In the time-analytical scheme they vary linearly in time
inside the step.

Both factories close over the `GridFields`, compile with `jax.jit`, and
return

```python
step_all(fluxes, state, intrpg, dsmin) -> ParticleState
```

which advances every active particle by one sub-cycle. The per-particle
work is done by `_subcycle_one_particle` (or its analytical twin) mapped
over the particles with `jax.vmap`. The patched and streaming integrators
use the same equations through `piecewise_local` and `piecewise_stencil`
(see [Integrators](integrators.md)).

Each sub-cycle works like `update_traj`, `cross_time` and `calc_pos` in the
Fortran code:

1. Interpolate the face fluxes in time.
2. Compute, for each of the six faces, the time the particle needs to reach
   it (`cross_time`).
3. Take the smallest of those times and the step cap `dsmin`.
4. Move the particle for that time (`calc_pos`) and, if it reached a face,
   move it into the neighbouring cell (`update_traj`).
5. Stop it if it ended up on land, outside the domain, or in a cell with no
   volume.

## Volume-normalised time

The schemes work with the time-like variable

$$
s = \frac{t}{\Delta V}, \qquad \Delta V = \Delta x\,\Delta y\,\Delta z ,
$$

so that a volume flux $U$ (m³ s⁻¹) in index space moves a particle by
$\mathrm{d}r/\mathrm{d}s = U$ cells. In the code $\Delta V$ is
`dxyz = dxdy * dzt * zstot`, interpolated in time with each particle's own
weight. A step `ds` corresponds to a physical time `dt = ds * dxyz`. The
step cap `dsmin`, the time left until the particle's next sub-step boundary,
is converted to `dsmin / dxyz` before the crossing times are compared.

## Time interpolation

A window lies between two GCM snapshots $t_0$ (slot 0) and $t_1$ (slot 1).
The fluxes on each face are interpolated linearly in time:

$$
U_i(t) = \frac{(t-t_0)\,U_i(t_1) + (t_1-t)\,U_i(t_0)}{t_1-t_0}
       = \texttt{intrpg}\,U_i(t_1) + (1-\texttt{intrpg})\,U_i(t_0).
$$

The weight `intrpg` is the fractional part of the normalised time `ts`
(time in units of the GCM timestep). Like the Fortran code, the
time-stepping scheme uses **one** weight for all particles in a sub-cycle:
the mean `ts` of the active particles. A per-particle weight lets the
interpolated flux at a wall change sign between consecutive sub-cycles,
which makes particles bounce back and forth across the wall. The cell
volume `dxyz` uses each particle's own weight, because it only affects the
step cap.

```{image} ../_static/window_timeline.svg
:class: diagram
:alt: A window between two snapshots divided into iter_sub sub-steps
:width: 720px
:align: center
```

## Time-stepping scheme (`"tstep"`)

Within a sub-step the fluxes are frozen at their time-interpolated values.
Consider the zonal direction in cell `ia`, which covers
$x_W = i_a - 1 \le x \le x_E = i_a$, with flux $U_W$ on the west face
(`um`) and $U_E$ on the east face (`uu`). TRACMASS assumes the flux varies
linearly across the cell:

$$
U(x) = U_W + (x - x_W)\,(U_E - U_W), \qquad x_W \le x \le x_E .
$$

With $\mathrm{d}x/\mathrm{d}s = U(x)$, this gives an exponential solution in
$s$.

### Crossing time (`_cross_time_1d`, Fortran `cross_time`)

For a particle at $x_0$, the time to reach the **east** face, `dse`, exists
when $U_E > 0$:

$$
s_E =
\begin{cases}
\dfrac{x_E - x_0}{U_E} & U_E = U_W, \\[2ex]
\dfrac{1}{U_W - U_E}\,\ln\!\left(\dfrac{U(x_0)}{U_E}\right) & U_E \ne U_W \text{ and } U(x_0) > 0,
\end{cases}
$$

and it is `UNDEF` ($10^{20}$) otherwise. The time to reach the **west**
face, `dsw`, exists when $U_W < 0$:

$$
s_W =
\begin{cases}
\dfrac{x_W - x_0}{U_E} & U_E = U_W, \\[2ex]
\dfrac{1}{U_W - U_E}\,\ln\!\left(\dfrac{U(x_0)}{U_W}\right) & U_E \ne U_W \text{ and } U(x_0) < 0 .
\end{cases}
$$

Non-positive results are set to `UNDEF`. The same function gives `dsn` and
`dss` from `vflux`, and `dsu` and `dsd` from `wflux`.

### Step length

The step is the smallest of the six crossing times and the cap:

$$
\Delta s = \min\left(s_E,\ s_W,\ s_N,\ s_S,\ s_U,\ s_D,\ \Delta s_{\min}\right).
$$

The face that sets $\Delta s$ decides where the particle goes, as in the
Fortran `update_traj`:

| Limiting time | New cell | Boxface |
|---|---|---|
| `dse` | `ia + 1` (wrapping to 1 with `iperio = 1`) | 1 east |
| `dsw` | `ia - 1` | 2 west |
| `dsn` | `ja + 1` | 3 north |
| `dss` | `ja - 1` | 4 south |
| `dsu` | `ka + 1` | 5 |
| `dsd` | `ka - 1` | 6 |
| `dsmin` | unchanged | 0, still inside the box |

If every crossing time is `UNDEF` (a convergence zone) the particle stays
in its box for the whole step. When two faces are reached at the same time
(a corner), the particle moves diagonally only if the diagonal cell is
ocean.

### New position (`_calc_pos_1d`, Fortran `calc_pos`)

After a step $\Delta s$ the new position is

$$
x_1 =
\begin{cases}
x_0 + U_E\,\Delta s & U_E = U_W, \\[1ex]
\left(x_0 - x_W + \dfrac{U_W}{U_E - U_W}\right) e^{(U_E - U_W)\,\Delta s}
  + x_W - \dfrac{U_W}{U_E - U_W} & U_E \ne U_W,
\end{cases}
$$

and the same for $y$ and $z$. On the face that set the step, the position
is set exactly to the wall.

:::{warning}
A particle at the exact point where $U(x) = 0$ in a divergent field is in
an unstable equilibrium. Neither the crossing time nor the new position is
defined there, as in the Fortran code.
:::

### Differences from the Fortran implementation

The equations are the same. A few safeguards were added for a vectorised
kernel running in single precision on large reanalysis grids. Each one is
covered by a test or a changelog entry:

- **Bounce at walls** (`update_bounce`). A particle sitting exactly on a
  wall, with the flux pushing it into the neighbouring cell, is moved into
  that cell before its crossing times are computed. The fluxes are then read
  again at the new indices.
- **No flow towards land.** Fluxes on faces next to a land cell
  (`kmt == 0`) are set to zero in the kernel, even if the data has a small
  residual there.
- **Wall nudge.** A position that lands exactly on a wall without crossing
  it is moved $10^{-5}$ cells back into the cell, which breaks endless
  corner oscillations.
- **Sea surface.** At `ka = 1` the flux through the outer vertical face is
  zero.
- **Domain edges.** On axes that do not wrap, the outermost ring of cells
  counts as outside the domain (see
  [Stopping trajectories](../tutorials/stopping.md#leaving-the-domain)).

### Sub-steps

`iter_sub` splits each window into equal sub-steps (Fortran `iter`). The cap
`dsmin` is the time left to the particle's next sub-step boundary, so no
step crosses one. Every particle therefore stops exactly on each sub-step
boundary, and so on every window boundary, where `tt` is a whole number of
timesteps. Stopping refreshes the time-interpolation weight, so larger
`iter_sub` values follow the time variation of the fluxes more closely,
at the cost of more sub-cycles.

## Time-analytical scheme (`"analytical"`)

The time-analytical scheme (Vries and Döös, 2001; Döös et al., 2017) lets
the fluxes vary linearly in time **and** space inside each step, and solves
the resulting equation exactly. It ignores `iter_sub` and integrates with
the GCM timestep `tseas`. The JAX implementation follows
`mod_pos_tanalytical.F90`.

With the face fluxes at the two snapshots

$$
U(x_E, t_0) = \texttt{uu}, \quad U(x_W, t_0) = \texttt{um}, \quad
U(x_E, t_1) = \texttt{vv}, \quad U(x_W, t_1) = \texttt{vm},
$$

the position $r = x/\Delta x$ obeys a linear equation in $r$ and $s$, with
coefficients

$$
\begin{aligned}
\alpha &= -\frac{1}{\Delta s}\,(\texttt{vv} - \texttt{vm} - \texttt{uu} + \texttt{um}), \\
\beta  &= (\texttt{um} - \texttt{uu}) - \alpha\,s_m, \\
\gamma &= -\frac{1}{\Delta s}\,(\texttt{vm} - \texttt{um}) - \alpha\,r_W, \\
\delta &= -\texttt{um} + r_W\,(\texttt{uu} - \texttt{um}) - \gamma\,s_m ,
\end{aligned}
$$

where $\Delta s$ is the volume-normalised GCM timestep, $s_m$ the time of
the past snapshot, and $r_W$ the index of the west wall. The solution
depends on the sign of $\alpha$:

- $\alpha > 0$, with $\xi = (\beta + \alpha s)/\sqrt{2\alpha}$ and Dawson's
  integral $D$:

  $$
  r(s) = \left(r_0 + \frac{\gamma}{\alpha}\right) e^{\xi_0^2-\xi^2} - \frac{\gamma}{\alpha}
       + \frac{\beta\gamma - \alpha\delta}{\alpha}\sqrt{\frac{2}{\alpha}}
         \left[D(\xi) - e^{\xi_0^2-\xi^2} D(\xi_0)\right]
  $$

- $\alpha < 0$, with $\zeta = (\beta + \alpha s)/\sqrt{-2\alpha}$ and the
  error function:

  $$
  r(s) = \left(r_0 + \frac{\gamma}{\alpha}\right) e^{\zeta^2-\zeta_0^2} - \frac{\gamma}{\alpha}
       - \frac{\beta\gamma - \alpha\delta}{\alpha}\sqrt{\frac{\pi}{-2\alpha}}\,
         e^{\zeta^2}\left[\operatorname{erf}(\zeta) - \operatorname{erf}(\zeta_0)\right]
  $$

- $\alpha = 0$ (rare in realistic fields):

  $$
  r(s) = \left(r_0 + \frac{\delta}{\beta}\right) e^{-\beta(s-s_0)} - \frac{\delta}{\beta}
       + \frac{\gamma}{\beta^2}\left[1 - \beta s + (\beta s_0 - 1)\,e^{-\beta(s-s_0)}\right]
  $$

The crossing time through a wall $r_w$ is the root of $r(s_w) - r_w = 0$.
Depending on whether the flux on the wall stays positive, changes sign or
stays negative between $t_0$ and $t_1$, there may be zero, one or two
candidate crossings. The Fortran code finds the root by bisection, in the
subroutines `apos`, `aneg` and `azer`. In JAX these are `_apos_crossing`,
`_aneg_crossing` and `_azer_crossing`, each with a fixed 100-step bisection
in `lax.fori_loop`, so that every particle runs the same number of
operations under `vmap`. Dawson's integral (`_dawson`) is a port of the
Fortran function `daw`: a series for $|x| < 0.2$ and Rybicki's sampling
method otherwise. The term $e^{\zeta^2}[\operatorname{erf}(\zeta) -
\operatorname{erf}(\zeta_0)]$ is computed by `_erfexp_val`, a port of
`erfexp`, which switches to an asymptotic expansion for $|\zeta| > 25$ to
avoid overflow.

The analytical scheme is available in `steady_state` and `time_dependent`.
It does not support ramp-up saves, because it cannot stop particles inside
a window. The patched and streaming integrators always use the
time-stepping scheme.

## Which scheme?

- Use `"tstep"` (the default) for most work. It is faster, it is supported
  by every integrator, and with `iter_sub` of 4 or more it follows
  time-varying fields well.
- Use `"analytical"` to compare with Fortran runs made with the
  time-analytical scheme, or when the exact solution of the linear-in-time
  field matters more than speed.

See Döös et al. (2017) for a comparison of the schemes with each other and
with Runge–Kutta integration.