Skip to content

Forcing Operators

Surface wind stress, bottom drag, and Rayleigh damping for C-grid ocean models.

Functional primitives

Pure, stateless functions that compute the core forcing math. They take fields already at the correct stagger and do no interpolation, ghost-ring padding, or masking. The output has the shape of the inputs, so to get a full [Ny, Nx] field slice every array input to the interior [1:-1, 1:-1] first, then pad the result back with interior() and apply any mask.

finitevolx.wind_stress_tendency(tau_on_face, rho0, dz_top)

Momentum tendency from a surface wind stress.

Converts a dynamic surface stress [N/m^2] (not the kinematic stress tau / rho0) into an acceleration of the top layer:

du_wind = tau / (rho0 * dz_top)

Parameters:

Name Type Description Default
tau_on_face Float[Array, '...']

Wind-stress component already interpolated to the target face points (U-points for tau_x, V-points for tau_y) [N/m^2].

 required
rho0 float or Float[Array, '...']

Reference density [kg/m^3].

 required
dz_top float or Float[Array, '...']

Top-layer thickness at the same face points [m]. Scalar or a field broadcastable to tau_on_face. Must be non-zero everywhere — replace dry-cell zeros before calling (e.g. jnp.where(mask, dz_top, 1.0)).

 required

Returns:

Type Description
Float[Array, '...']

Wind-stress tendency [m/s^2], same shape as tau_on_face.

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import wind_stress_tendency
>>> du = wind_stress_tendency(jnp.array([1025.0]), rho0=1025.0, dz_top=2.0)
>>> float(du[0])
0.5
Source code in finitevolx/_src/forcing/functional.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
def wind_stress_tendency(
    tau_on_face: Float[Array, "..."],
    rho0: float | Float[Array, "..."],
    dz_top: float | Float[Array, "..."],
) -> Float[Array, "..."]:
    """Momentum tendency from a surface wind stress.

    Converts a dynamic surface stress [N/m^2] (not the kinematic stress
    ``tau / rho0``) into an acceleration of the top layer:

        du_wind = tau / (rho0 * dz_top)

    Parameters
    ----------
    tau_on_face : Float[Array, "..."]
        Wind-stress component **already interpolated** to the target face
        points (U-points for ``tau_x``, V-points for ``tau_y``) [N/m^2].
    rho0 : float or Float[Array, "..."]
        Reference density [kg/m^3].
    dz_top : float or Float[Array, "..."]
        Top-layer thickness at the same face points [m].  Scalar or a
        field broadcastable to ``tau_on_face``.  Must be non-zero
        everywhere — replace dry-cell zeros before calling (e.g.
        ``jnp.where(mask, dz_top, 1.0)``).

    Returns
    -------
    Float[Array, "..."]
        Wind-stress tendency [m/s^2], same shape as ``tau_on_face``.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import wind_stress_tendency
    >>> du = wind_stress_tendency(jnp.array([1025.0]), rho0=1025.0, dz_top=2.0)
    >>> float(du[0])
    0.5
    """
    # du_wind = tau / (rho0 * dz_top)
    return tau_on_face / (rho0 * dz_top)

finitevolx.linear_drag_tendency(vel, r)

Linear (Rayleigh-type) drag tendency on one velocity component.

dvel = -r * vel

Parameters:

Name Type Description Default
vel Float[Array, '...']

Velocity component at its native face points (U or V).

 required
r float or Float[Array, '...']

Drag coefficient [1/s]. Scalar or a field at the same face points.

 required

Returns:

Type Description
Float[Array, '...']

Drag tendency [m/s^2], same shape as vel.

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import linear_drag_tendency
>>> linear_drag_tendency(jnp.array([2.0, -4.0]), r=0.5).tolist()
[-1.0, 2.0]
Source code in finitevolx/_src/forcing/functional.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
def linear_drag_tendency(
    vel: Float[Array, "..."],
    r: float | Float[Array, "..."],
) -> Float[Array, "..."]:
    """Linear (Rayleigh-type) drag tendency on one velocity component.

        dvel = -r * vel

    Parameters
    ----------
    vel : Float[Array, "..."]
        Velocity component at its native face points (U or V).
    r : float or Float[Array, "..."]
        Drag coefficient [1/s].  Scalar or a field at the same face points.

    Returns
    -------
    Float[Array, "..."]
        Drag tendency [m/s^2], same shape as ``vel``.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import linear_drag_tendency
    >>> linear_drag_tendency(jnp.array([2.0, -4.0]), r=0.5).tolist()
    [-1.0, 2.0]
    """
    # dvel = -r * vel
    return -r * vel

finitevolx.quadratic_drag_tendency(vel, speed, cd, h_bot)

Quadratic bottom-drag tendency on one velocity component.

dvel = -Cd * |u| * vel / h_bot

Parameters:

Name Type Description Default
vel Float[Array, '...']

Velocity component at its native face points (U or V).

 required
speed Float[Array, '...']

Flow speed |u| at the same face points [m/s].

 required
cd float or Float[Array, '...']

Dimensionless drag coefficient (typically ~1e-3).

 required
h_bot float or Float[Array, '...']

Bottom-layer thickness at the same face points [m]. Must be non-zero everywhere — replace dry-cell zeros before calling (e.g. jnp.where(mask, h_bot, 1.0)).

 required

Returns:

Type Description
Float[Array, '...']

Drag tendency [m/s^2], same shape as vel.

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import quadratic_drag_tendency
>>> vel = jnp.array([2.0])
>>> float(quadratic_drag_tendency(vel, jnp.abs(vel), cd=0.5, h_bot=4.0)[0])
-0.5
Source code in finitevolx/_src/forcing/functional.py 
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def quadratic_drag_tendency(
    vel: Float[Array, "..."],
    speed: Float[Array, "..."],
    cd: float | Float[Array, "..."],
    h_bot: float | Float[Array, "..."],
) -> Float[Array, "..."]:
    """Quadratic bottom-drag tendency on one velocity component.

        dvel = -Cd * |u| * vel / h_bot

    Parameters
    ----------
    vel : Float[Array, "..."]
        Velocity component at its native face points (U or V).
    speed : Float[Array, "..."]
        Flow speed ``|u|`` at the same face points [m/s].
    cd : float or Float[Array, "..."]
        Dimensionless drag coefficient (typically ``~1e-3``).
    h_bot : float or Float[Array, "..."]
        Bottom-layer thickness at the same face points [m].  Must be
        non-zero everywhere — replace dry-cell zeros before calling (e.g.
        ``jnp.where(mask, h_bot, 1.0)``).

    Returns
    -------
    Float[Array, "..."]
        Drag tendency [m/s^2], same shape as ``vel``.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import quadratic_drag_tendency
    >>> vel = jnp.array([2.0])
    >>> float(quadratic_drag_tendency(vel, jnp.abs(vel), cd=0.5, h_bot=4.0)[0])
    -0.5
    """
    # dvel = -Cd * |u| * vel / h_bot
    return -cd * speed * vel / h_bot

finitevolx.rayleigh_tendency(q, r, q_ref=None)

Rayleigh damping tendency toward a reference state.

dq = -r * (q - q_ref)

Parameters:

Name Type Description Default
q Float[Array, '...']

Field value at its native grid points.

 required
r float or Float[Array, '...']

Damping rate [1/s]. Scalar or a field at the same points.

 required
q_ref (float, Float[Array, '...'] or None)

Reference state broadcastable to q. None (default) means pure linear damping toward zero.

 None

Returns:

Type Description
Float[Array, '...']

Damping tendency, same shape as q.

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import rayleigh_tendency
>>> q = jnp.array([3.0])
>>> float(rayleigh_tendency(q, r=0.5)[0])
-1.5
>>> float(rayleigh_tendency(q, r=0.5, q_ref=1.0)[0])
-1.0
Source code in finitevolx/_src/forcing/functional.py 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
def rayleigh_tendency(
    q: Float[Array, "..."],
    r: float | Float[Array, "..."],
    q_ref: float | Float[Array, "..."] | None = None,
) -> Float[Array, "..."]:
    """Rayleigh damping tendency toward a reference state.

        dq = -r * (q - q_ref)

    Parameters
    ----------
    q : Float[Array, "..."]
        Field value at its native grid points.
    r : float or Float[Array, "..."]
        Damping rate [1/s].  Scalar or a field at the same points.
    q_ref : float, Float[Array, "..."] or None, optional
        Reference state broadcastable to ``q``.  ``None`` (default) means
        pure linear damping toward zero.

    Returns
    -------
    Float[Array, "..."]
        Damping tendency, same shape as ``q``.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import rayleigh_tendency
    >>> q = jnp.array([3.0])
    >>> float(rayleigh_tendency(q, r=0.5)[0])
    -1.5
    >>> float(rayleigh_tendency(q, r=0.5, q_ref=1.0)[0])
    -1.0
    """
    if q_ref is None:
        # dq = -r * q
        return -r * q
    # dq = -r * (q - q_ref)
    return -r * (q - q_ref)

Base class and composition

AbstractForcing is the contract every forcing operator satisfies. The composition wrappers combine forcings inside a model's right-hand side: ForcingSum adds independent processes, ForcingProduct modulates a forcing by a fixed or callable factor (sponges, seasonal cycles), and TimeVaryingForcing evaluates time-dependent input fields — analytic or interpolated with diffrax.LinearInterpolation — before calling a time-agnostic forcing.

finitevolx.AbstractForcing

Bases: Module

Base class for forcing operators.

Subclasses implement __call__ to return a tendency contribution — a single array for scalar forcings (e.g. dq) or a tuple of arrays for momentum forcings (e.g. (du, dv)). The call signature is left to each subclass.

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import AbstractForcing
>>> class ConstantForcing(AbstractForcing):
...     value: float
...
...     def __call__(self, q):
...         return self.value * jnp.ones_like(q)
>>> ConstantForcing(value=2.0)(jnp.zeros(3)).tolist()
[2.0, 2.0, 2.0]
Source code in finitevolx/_src/forcing/_base.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
class AbstractForcing(eqx.Module):
    """Base class for forcing operators.

    Subclasses implement ``__call__`` to return a tendency contribution —
    a single array for scalar forcings (e.g. ``dq``) or a tuple of arrays for
    momentum forcings (e.g. ``(du, dv)``).  The call signature is left to
    each subclass.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import AbstractForcing
    >>> class ConstantForcing(AbstractForcing):
    ...     value: float
    ...
    ...     def __call__(self, q):
    ...         return self.value * jnp.ones_like(q)
    >>> ConstantForcing(value=2.0)(jnp.zeros(3)).tolist()
    [2.0, 2.0, 2.0]
    """

    @abc.abstractmethod
    def __call__(self, *args: Any, **kwargs: Any) -> PyTree[Array]:
        """Return the tendency contribution of this forcing."""
        raise NotImplementedError

__call__(*args, **kwargs) abstractmethod

Return the tendency contribution of this forcing.

Source code in finitevolx/_src/forcing/_base.py 
56
57
58
59
@abc.abstractmethod
def __call__(self, *args: Any, **kwargs: Any) -> PyTree[Array]:
    """Return the tendency contribution of this forcing."""
    raise NotImplementedError

finitevolx.ForcingSum

Bases: AbstractForcing

Additive composition of forcings: F = F_1 + F_2 + ... + F_n.

Every member is called with the same arguments and the resulting tendencies are summed leaf-by-leaf, so all members must return pytrees of the same structure (e.g. all (du, dv) tuples). Analogous to diffrax.MultiTerm for the pieces of a single vector field.

Parameters:

Name Type Description Default
*forcings AbstractForcing

One or more forcings sharing a call signature and output structure.

 ()

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import AbstractForcing, ForcingSum
>>> class Scale(AbstractForcing):
...     a: float
...
...     def __call__(self, q):
...         return self.a * q
>>> total = ForcingSum(Scale(1.0), Scale(2.0))
>>> total(jnp.ones(2)).tolist()
[3.0, 3.0]
Source code in finitevolx/_src/forcing/_base.py 
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
class ForcingSum(AbstractForcing):
    """Additive composition of forcings: ``F = F_1 + F_2 + ... + F_n``.

    Every member is called with the same arguments and the resulting
    tendencies are summed leaf-by-leaf, so all members must return pytrees
    of the same structure (e.g. all ``(du, dv)`` tuples).  Analogous to
    ``diffrax.MultiTerm`` for the pieces of a single vector field.

    Parameters
    ----------
    *forcings : AbstractForcing
        One or more forcings sharing a call signature and output structure.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import AbstractForcing, ForcingSum
    >>> class Scale(AbstractForcing):
    ...     a: float
    ...
    ...     def __call__(self, q):
    ...         return self.a * q
    >>> total = ForcingSum(Scale(1.0), Scale(2.0))
    >>> total(jnp.ones(2)).tolist()
    [3.0, 3.0]
    """

    forcings: tuple[AbstractForcing, ...]

    def __init__(self, *forcings: AbstractForcing) -> None:
        if not forcings:
            raise ValueError("ForcingSum requires at least one forcing.")
        self.forcings = tuple(forcings)

    def __call__(self, *args: Any, **kwargs: Any) -> PyTree[Array]:
        """Sum of every member's tendency, computed leaf-by-leaf.

        Parameters
        ----------
        *args, **kwargs
            Forwarded unchanged to every member forcing.

        Returns
        -------
        PyTree[Array]
            Leaf-wise sum of the member tendencies.
        """
        results = [f(*args, **kwargs) for f in self.forcings]
        # F[leaf] = F_1[leaf] + F_2[leaf] + ... + F_n[leaf]
        return jax.tree.map(lambda *xs: functools.reduce(operator.add, xs), *results)

__call__(*args, **kwargs)

Sum of every member's tendency, computed leaf-by-leaf.

Parameters:

Name Type Description Default
*args Any

Forwarded unchanged to every member forcing.

 ()
**kwargs Any

Forwarded unchanged to every member forcing.

 ()

Returns:

Type Description
PyTree[Array]

Leaf-wise sum of the member tendencies.
Source code in finitevolx/_src/forcing/_base.py 
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
def __call__(self, *args: Any, **kwargs: Any) -> PyTree[Array]:
    """Sum of every member's tendency, computed leaf-by-leaf.

    Parameters
    ----------
    *args, **kwargs
        Forwarded unchanged to every member forcing.

    Returns
    -------
    PyTree[Array]
        Leaf-wise sum of the member tendencies.
    """
    results = [f(*args, **kwargs) for f in self.forcings]
    # F[leaf] = F_1[leaf] + F_2[leaf] + ... + F_n[leaf]
    return jax.tree.map(lambda *xs: functools.reduce(operator.add, xs), *results)

finitevolx.ForcingProduct

Bases: AbstractForcing

Modulated forcing: F = alpha * F_inner.

Scales every leaf of the inner forcing's tendency by alpha. Typical uses are sponge layers (alpha(x) a smooth 0 -> 1 ramp near a boundary), seasonal modulation (alpha(t)), or a learned efficiency field (alpha a trainable array).

Parameters:

Name Type Description Default
forcing AbstractForcing

The inner forcing.

 required
modulator Callable or float or Float[Array, '...']

Either a fixed factor (scalar or array broadcastable to each output leaf) or a callable. A callable modulator is called with the same arguments as the forcing and must return the factor.

 required

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import AbstractForcing, ForcingProduct
>>> class Identity(AbstractForcing):
...     def __call__(self, q):
...         return q
>>> sponge = jnp.array([0.0, 0.5, 1.0])
>>> ForcingProduct(Identity(), sponge)(jnp.full(3, 2.0)).tolist()
[0.0, 1.0, 2.0]
Source code in finitevolx/_src/forcing/_base.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
class ForcingProduct(AbstractForcing):
    """Modulated forcing: ``F = alpha * F_inner``.

    Scales every leaf of the inner forcing's tendency by ``alpha``.  Typical
    uses are sponge layers (``alpha(x)`` a smooth 0 -> 1 ramp near a
    boundary), seasonal modulation (``alpha(t)``), or a learned efficiency
    field (``alpha`` a trainable array).

    Parameters
    ----------
    forcing : AbstractForcing
        The inner forcing.
    modulator : Callable or float or Float[Array, "..."]
        Either a fixed factor (scalar or array broadcastable to each output
        leaf) or a callable.  A callable modulator is called with the **same
        arguments** as the forcing and must return the factor.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import AbstractForcing, ForcingProduct
    >>> class Identity(AbstractForcing):
    ...     def __call__(self, q):
    ...         return q
    >>> sponge = jnp.array([0.0, 0.5, 1.0])
    >>> ForcingProduct(Identity(), sponge)(jnp.full(3, 2.0)).tolist()
    [0.0, 1.0, 2.0]
    """

    forcing: AbstractForcing
    modulator: Callable[..., Any] | float | Float[Array, "..."]

    def __call__(self, *args: Any, **kwargs: Any) -> PyTree[Array]:
        """Modulated tendency ``alpha * F_inner(*args, **kwargs)``.

        Parameters
        ----------
        *args, **kwargs
            Forwarded unchanged to the inner forcing (and to the modulator
            when it is callable).

        Returns
        -------
        PyTree[Array]
            Inner tendency with every leaf multiplied by ``alpha``.
        """
        tendency = self.forcing(*args, **kwargs)
        modulator = self.modulator
        alpha = (
            cast(Callable[..., Any], modulator)(*args, **kwargs)
            if callable(modulator)
            else modulator
        )
        # F[leaf] = alpha * F_inner[leaf]
        return jax.tree.map(lambda t: alpha * t, tendency)

__call__(*args, **kwargs)

Modulated tendency alpha * F_inner(*args, **kwargs).

Parameters:

Name Type Description Default
*args Any

Forwarded unchanged to the inner forcing (and to the modulator when it is callable).

 ()
**kwargs Any

Forwarded unchanged to the inner forcing (and to the modulator when it is callable).

 ()

Returns:

Type Description
PyTree[Array]

Inner tendency with every leaf multiplied by alpha.
Source code in finitevolx/_src/forcing/_base.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
def __call__(self, *args: Any, **kwargs: Any) -> PyTree[Array]:
    """Modulated tendency ``alpha * F_inner(*args, **kwargs)``.

    Parameters
    ----------
    *args, **kwargs
        Forwarded unchanged to the inner forcing (and to the modulator
        when it is callable).

    Returns
    -------
    PyTree[Array]
        Inner tendency with every leaf multiplied by ``alpha``.
    """
    tendency = self.forcing(*args, **kwargs)
    modulator = self.modulator
    alpha = (
        cast(Callable[..., Any], modulator)(*args, **kwargs)
        if callable(modulator)
        else modulator
    )
    # F[leaf] = alpha * F_inner[leaf]
    return jax.tree.map(lambda t: alpha * t, tendency)

finitevolx.TimeVaryingForcing

Bases: AbstractForcing

Bundle a time-agnostic forcing with time-dependent input fields.

At call time the input fields are evaluated as field_fn(t) and passed to forcing as its leading positional arguments; a tuple is unpacked into several arguments, anything else is passed as a single argument. Any further *args / **kwargs follow the evaluated fields.

field_fn may be an analytic formula or a data-driven interpolant such as diffrax.LinearInterpolation(ts, ys).evaluate (ERA5, JRA-55, ...), both of which are compatible with jax.jit / jax.grad.

Parameters:

Name Type Description Default
forcing Callable

Any forcing operator (typically a Layer 2 operator).

 required
field_fn Callable

t -> fields — a single array or a tuple of arrays.

 required

Examples:

Analytic seasonal modulation of a scalar forcing:

>>> import jax.numpy as jnp
>>> from finitevolx import AbstractForcing, TimeVaryingForcing
>>> class Scale(AbstractForcing):
...     def __call__(self, q, a):
...         return a * q
>>> tv = TimeVaryingForcing(Scale(), lambda t: jnp.full(2, 1.0 + t))
>>> tv(1.0, a=3.0).tolist()
[6.0, 6.0]

Data-driven wind from gridded reanalysis (sketch)::

import diffrax as dfx

tau_x_path = dfx.LinearInterpolation(ts=era5_times, ys=era5_tau_x)
tau_y_path = dfx.LinearInterpolation(ts=era5_times, ys=era5_tau_y)
tv_wind = TimeVaryingForcing(
    forcing=wind_op,
    field_fn=lambda t: (tau_x_path.evaluate(t), tau_y_path.evaluate(t)),
)
du, dv = tv_wind(t, dz_top=50.0)
Source code in finitevolx/_src/forcing/_base.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
class TimeVaryingForcing(AbstractForcing):
    """Bundle a time-agnostic forcing with time-dependent input fields.

    At call time the input fields are evaluated as ``field_fn(t)`` and passed
    to ``forcing`` as its leading positional arguments; a tuple is unpacked
    into several arguments, anything else is passed as a single argument.
    Any further ``*args`` / ``**kwargs`` follow the evaluated fields.

    ``field_fn`` may be an analytic formula or a data-driven interpolant such
    as ``diffrax.LinearInterpolation(ts, ys).evaluate`` (ERA5, JRA-55, ...),
    both of which are compatible with ``jax.jit`` / ``jax.grad``.

    Parameters
    ----------
    forcing : Callable
        Any forcing operator (typically a Layer 2 operator).
    field_fn : Callable
        ``t -> fields`` — a single array or a tuple of arrays.

    Examples
    --------
    Analytic seasonal modulation of a scalar forcing:

    >>> import jax.numpy as jnp
    >>> from finitevolx import AbstractForcing, TimeVaryingForcing
    >>> class Scale(AbstractForcing):
    ...     def __call__(self, q, a):
    ...         return a * q
    >>> tv = TimeVaryingForcing(Scale(), lambda t: jnp.full(2, 1.0 + t))
    >>> tv(1.0, a=3.0).tolist()
    [6.0, 6.0]

    Data-driven wind from gridded reanalysis (sketch)::

        import diffrax as dfx

        tau_x_path = dfx.LinearInterpolation(ts=era5_times, ys=era5_tau_x)
        tau_y_path = dfx.LinearInterpolation(ts=era5_times, ys=era5_tau_y)
        tv_wind = TimeVaryingForcing(
            forcing=wind_op,
            field_fn=lambda t: (tau_x_path.evaluate(t), tau_y_path.evaluate(t)),
        )
        du, dv = tv_wind(t, dz_top=50.0)
    """

    forcing: Callable[..., PyTree[Array]]
    field_fn: Callable[[Any], Any]

    def __call__(self, t: Any, *args: Any, **kwargs: Any) -> PyTree[Array]:
        """Evaluate the fields at ``t`` and apply the wrapped forcing.

        Parameters
        ----------
        t : scalar
            Time at which to evaluate ``field_fn``.
        *args, **kwargs
            Extra arguments forwarded to the forcing after the fields.

        Returns
        -------
        PyTree[Array]
            The wrapped forcing's tendency.
        """
        fields = self.field_fn(t)
        if isinstance(fields, tuple):
            return self.forcing(*fields, *args, **kwargs)
        return self.forcing(fields, *args, **kwargs)

__call__(t, *args, **kwargs)

Evaluate the fields at t and apply the wrapped forcing.

Parameters:

Name Type Description Default
t scalar

Time at which to evaluate field_fn.

 required
*args Any

Extra arguments forwarded to the forcing after the fields.

 ()
**kwargs Any

Extra arguments forwarded to the forcing after the fields.

 ()

Returns:

Type Description
PyTree[Array]

The wrapped forcing's tendency.
Source code in finitevolx/_src/forcing/_base.py 
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
def __call__(self, t: Any, *args: Any, **kwargs: Any) -> PyTree[Array]:
    """Evaluate the fields at ``t`` and apply the wrapped forcing.

    Parameters
    ----------
    t : scalar
        Time at which to evaluate ``field_fn``.
    *args, **kwargs
        Extra arguments forwarded to the forcing after the fields.

    Returns
    -------
    PyTree[Array]
        The wrapped forcing's tendency.
    """
    fields = self.field_fn(t)
    if isinstance(fields, tuple):
        return self.forcing(*fields, *args, **kwargs)
    return self.forcing(fields, *args, **kwargs)

2D operators

Module operators for CartesianGrid2D. They interpolate T-point inputs to the velocity faces, delegate the math to the functional primitives, restore the zero ghost ring, and apply an optional Mask2D given at construction. Multilayer models ([nl, Ny, Nx]) reuse them through multilayer() or by calling them on a single layer, e.g. bottom drag on u[-1], v[-1].

finitevolx.WindStress2D

Bases: AbstractForcing

Surface wind-stress forcing on a 2-D Arakawa C-grid.

Computes the momentum tendency of the top layer:

du_wind[j, i+1/2] = tau_x_on_u[j, i+1/2] / (rho0 * dz_top_on_u[j, i+1/2])
dv_wind[j+1/2, i] = tau_y_on_v[j+1/2, i] / (rho0 * dz_top_on_v[j+1/2, i])

tau_x, tau_y (and dz_top when it is a field) are given at T-points and interpolated to U/V-points with simple x/y averaging.

Parameters:

Name Type Description Default
grid CartesianGrid2D

The underlying 2-D grid.

 required
mask Mask2D or None

Optional land/ocean mask. When provided, du_wind is zeroed at dry U-faces and dv_wind at dry V-faces (via jnp.where, so a zero dry-cell dz_top or NaN-filled land never leaks NaN).

 None
rho0 float

Reference density [kg/m^3]. Default 1025.0.

 1025.0

Examples:

Idealised double-gyre wind tau_x = -tau0 * cos(2 pi y / Ly):

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid2D, WindStress2D
>>> grid = CartesianGrid2D.from_interior(8, 8, 1.0e6, 1.0e6)
>>> y = (jnp.arange(grid.Ny) - 0.5) * grid.dy
>>> tau_x = -0.1 * jnp.cos(2 * jnp.pi * y / grid.Ly)[:, None] * jnp.ones(grid.Nx)
>>> tau_y = jnp.zeros((grid.Ny, grid.Nx))
>>> wind = WindStress2D(grid=grid)
>>> du_wind, dv_wind = wind(tau_x, tau_y, dz_top=50.0)
Source code in finitevolx/_src/forcing/wind.py 
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
class WindStress2D(AbstractForcing):
    """Surface wind-stress forcing on a 2-D Arakawa C-grid.

    Computes the momentum tendency of the top layer:

        du_wind[j, i+1/2] = tau_x_on_u[j, i+1/2] / (rho0 * dz_top_on_u[j, i+1/2])
        dv_wind[j+1/2, i] = tau_y_on_v[j+1/2, i] / (rho0 * dz_top_on_v[j+1/2, i])

    ``tau_x``, ``tau_y`` (and ``dz_top`` when it is a field) are given at
    T-points and interpolated to U/V-points with simple x/y averaging.

    Parameters
    ----------
    grid : CartesianGrid2D
        The underlying 2-D grid.
    mask : Mask2D or None, optional
        Optional land/ocean mask.  When provided, ``du_wind`` is zeroed at
        dry U-faces and ``dv_wind`` at dry V-faces (via ``jnp.where``, so a
        zero dry-cell ``dz_top`` or NaN-filled land never leaks NaN).
    rho0 : float, optional
        Reference density [kg/m^3].  Default ``1025.0``.

    Examples
    --------
    Idealised double-gyre wind ``tau_x = -tau0 * cos(2 pi y / Ly)``:

    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid2D, WindStress2D
    >>> grid = CartesianGrid2D.from_interior(8, 8, 1.0e6, 1.0e6)
    >>> y = (jnp.arange(grid.Ny) - 0.5) * grid.dy
    >>> tau_x = -0.1 * jnp.cos(2 * jnp.pi * y / grid.Ly)[:, None] * jnp.ones(grid.Nx)
    >>> tau_y = jnp.zeros((grid.Ny, grid.Nx))
    >>> wind = WindStress2D(grid=grid)
    >>> du_wind, dv_wind = wind(tau_x, tau_y, dz_top=50.0)
    """

    grid: CartesianGrid2D
    mask: Mask2D | None
    rho0: float
    interp: Interpolation2D

    def __init__(
        self,
        grid: CartesianGrid2D,
        mask: Mask2D | None = None,
        rho0: float = 1025.0,
    ) -> None:
        self.grid = grid
        self.mask = mask
        self.rho0 = rho0
        self.interp = Interpolation2D(grid=grid)

    def __call__(
        self,
        tau_x: Float[Array, "Ny Nx"],
        tau_y: Float[Array, "Ny Nx"],
        dz_top: float | Float[Array, "Ny Nx"],
    ) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
        """Wind-stress tendencies (du_wind, dv_wind).

        Parameters
        ----------
        tau_x : Float[Array, "Ny Nx"]
            Zonal wind stress at T-points [N/m^2].
        tau_y : Float[Array, "Ny Nx"]
            Meridional wind stress at T-points [N/m^2].
        dz_top : float or Float[Array, "Ny Nx"]
            Top-layer thickness [m] — a scalar, or a T-point field (ghost
            ring filled) interpolated to U/V-points.

        Returns
        -------
        tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]
            ``(du_wind, dv_wind)`` at U- and V-points, zero in the ghost
            ring.  When ``self.mask`` is set, dry faces are zeroed.
        """
        mu = None if self.mask is None else self.mask.u
        mv = None if self.mask is None else self.mask.v
        return self._tendency(tau_x, tau_y, dz_top, mu, mv)

    def _tendency(
        self,
        tau_x: Float[Array, "Ny Nx"],
        tau_y: Float[Array, "Ny Nx"],
        dz_top: float | Float[Array, "Ny Nx"],
        mu: Array | None,
        mv: Array | None,
    ) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
        """Core stencil with explicit ``[Ny, Nx]`` face masks (or ``None``).

        Shared with :class:`WindStress3D`, which passes one z-level's masks.
        """
        # tau_x_on_u[j, i+1/2] = 1/2 * (tau_x[j, i] + tau_x[j, i+1])
        # (zeroed at dry faces so NaN land stress cannot give NaN gradients)
        tau_x_on_u = mask_where(self.interp.T_to_U(tau_x), mu)
        # tau_y_on_v[j+1/2, i] = 1/2 * (tau_y[j, i] + tau_y[j+1, i])
        tau_y_on_v = mask_where(self.interp.T_to_V(tau_y), mv)
        # Swap in 1 at dry faces so a zero dry-cell thickness cannot give 0/0.
        dz_on_u = safe_denominator(on_face_interior(dz_top, self.interp.T_to_U), mu)
        dz_on_v = safe_denominator(on_face_interior(dz_top, self.interp.T_to_V), mv)

        # du_wind[j, i+1/2] = tau_x_on_u / (rho0 * dz_top_on_u)
        du_wind = interior(
            wind_stress_tendency(tau_x_on_u[1:-1, 1:-1], self.rho0, dz_on_u), tau_x
        )
        # dv_wind[j+1/2, i] = tau_y_on_v / (rho0 * dz_top_on_v)
        dv_wind = interior(
            wind_stress_tendency(tau_y_on_v[1:-1, 1:-1], self.rho0, dz_on_v), tau_y
        )

        return mask_where(du_wind, mu), mask_where(dv_wind, mv)

__call__(tau_x, tau_y, dz_top)

Wind-stress tendencies (du_wind, dv_wind).

Parameters:

Name Type Description Default
tau_x Float[Array, 'Ny Nx']

Zonal wind stress at T-points [N/m^2].

 required
tau_y Float[Array, 'Ny Nx']

Meridional wind stress at T-points [N/m^2].

 required
dz_top float or Float[Array, 'Ny Nx']

Top-layer thickness [m] — a scalar, or a T-point field (ghost ring filled) interpolated to U/V-points.

 required

Returns:

Type Description
tuple[Float[Array, 'Ny Nx'], Float[Array, 'Ny Nx']]

(du_wind, dv_wind) at U- and V-points, zero in the ghost ring. When self.mask is set, dry faces are zeroed.
Source code in finitevolx/_src/forcing/wind.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def __call__(
    self,
    tau_x: Float[Array, "Ny Nx"],
    tau_y: Float[Array, "Ny Nx"],
    dz_top: float | Float[Array, "Ny Nx"],
) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
    """Wind-stress tendencies (du_wind, dv_wind).

    Parameters
    ----------
    tau_x : Float[Array, "Ny Nx"]
        Zonal wind stress at T-points [N/m^2].
    tau_y : Float[Array, "Ny Nx"]
        Meridional wind stress at T-points [N/m^2].
    dz_top : float or Float[Array, "Ny Nx"]
        Top-layer thickness [m] — a scalar, or a T-point field (ghost
        ring filled) interpolated to U/V-points.

    Returns
    -------
    tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]
        ``(du_wind, dv_wind)`` at U- and V-points, zero in the ghost
        ring.  When ``self.mask`` is set, dry faces are zeroed.
    """
    mu = None if self.mask is None else self.mask.u
    mv = None if self.mask is None else self.mask.v
    return self._tendency(tau_x, tau_y, dz_top, mu, mv)

WindCurl2D is the vorticity / QG-PV form of the same forcing: the C-grid curl of the face-interpolated stress divided by rho0 * H, at X-points.

finitevolx.WindCurl2D

Bases: AbstractForcing

Wind-stress curl forcing of the vorticity / QG potential vorticity.

The vorticity tendency of a layer of thickness H driven by the surface stress (tau_x, tau_y), at X-points (corners):

F[j+1/2, i+1/2] = ( (tau_y_on_v[j+1/2, i+1] - tau_y_on_v[j+1/2, i]) / dx
                  - (tau_x_on_u[j+1, i+1/2] - tau_x_on_u[j, i+1/2]) / dy )
                  / (rho0 * H)

i.e. the curl of the stress interpolated to the faces (the C-grid curl of :class:Difference2D / :class:Vorticity2D), divided by rho0 * H. tau_x, tau_y are given at T-points with the ghost ring filled (e.g. an analytic stress evaluated at the ghost-cell centres): the face averages are formed on the full arrays, so the X-points next to the north / east walls read the ghost stress. For the double-gyre tau_x = -tau0 cos(2 pi y / Ly) this gives F = -tau0 (2 pi / Ly) sin(2 pi y / Ly) / (rho0 H): negative (anticyclonic) in the southern half.

Parameters:

Name Type Description Default
grid CartesianGrid2D

The underlying 2-D grid.

 required
mask Mask2D or None

Optional land/ocean mask. Dry T-point stress is zeroed with jnp.where before the stencil (so NaN-filled land cannot leak) and the output is zeroed on dry X-points (mask.xy_corner_strict).

 required
rho0 float

Reference density [kg/m^3]. Default 1025.0.

 required

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid2D, WindCurl2D
>>> grid = CartesianGrid2D.from_interior(32, 32, 5.12e6, 5.12e6)
>>> y = (jnp.arange(grid.Ny) - 0.5) * grid.dy  # T-point y, ghosts included
>>> tau_x = -0.1 * jnp.cos(2 * jnp.pi * y / grid.Ly)[:, None] * jnp.ones(grid.Nx)
>>> curl = WindCurl2D(grid=grid)(tau_x, jnp.zeros_like(tau_x), H=500.0)
>>> curl.shape  # X-points, zero ghost ring
(34, 34)
Source code in finitevolx/_src/forcing/wind.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
class WindCurl2D(AbstractForcing):
    """Wind-stress curl forcing of the vorticity / QG potential vorticity.

    The vorticity tendency of a layer of thickness ``H`` driven by the
    surface stress ``(tau_x, tau_y)``, at X-points (corners):

        F[j+1/2, i+1/2] = ( (tau_y_on_v[j+1/2, i+1] - tau_y_on_v[j+1/2, i]) / dx
                          - (tau_x_on_u[j+1, i+1/2] - tau_x_on_u[j, i+1/2]) / dy )
                          / (rho0 * H)

    i.e. the curl of the stress interpolated to the faces (the C-grid curl of
    :class:`Difference2D` / :class:`Vorticity2D`), divided by ``rho0 * H``.
    ``tau_x``, ``tau_y`` are given at T-points **with the ghost ring filled**
    (e.g. an analytic stress evaluated at the ghost-cell centres): the face
    averages are formed on the full arrays, so the X-points next to the
    north / east walls read the ghost stress.  For the double-gyre
    ``tau_x = -tau0 cos(2 pi y / Ly)`` this gives
    ``F = -tau0 (2 pi / Ly) sin(2 pi y / Ly) / (rho0 H)``: negative
    (anticyclonic) in the southern half.

    Parameters
    ----------
    grid : CartesianGrid2D
        The underlying 2-D grid.
    mask : Mask2D or None, optional
        Optional land/ocean mask.  Dry T-point stress is zeroed with
        ``jnp.where`` before the stencil (so NaN-filled land cannot leak)
        and the output is zeroed on dry X-points (``mask.xy_corner_strict``).
    rho0 : float, optional
        Reference density [kg/m^3].  Default ``1025.0``.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid2D, WindCurl2D
    >>> grid = CartesianGrid2D.from_interior(32, 32, 5.12e6, 5.12e6)
    >>> y = (jnp.arange(grid.Ny) - 0.5) * grid.dy  # T-point y, ghosts included
    >>> tau_x = -0.1 * jnp.cos(2 * jnp.pi * y / grid.Ly)[:, None] * jnp.ones(grid.Nx)
    >>> curl = WindCurl2D(grid=grid)(tau_x, jnp.zeros_like(tau_x), H=500.0)
    >>> curl.shape  # X-points, zero ghost ring
    (34, 34)
    """

    grid: CartesianGrid2D
    mask: Mask2D | None = None
    rho0: float = 1025.0

    def __call__(
        self,
        tau_x: Float[Array, "Ny Nx"],
        tau_y: Float[Array, "Ny Nx"],
        H: float | Float[Array, "Ny Nx"],
    ) -> Float[Array, "Ny Nx"]:
        """Wind-curl forcing ``curl(tau) / (rho0 H)`` at X-points.

        Parameters
        ----------
        tau_x : Float[Array, "Ny Nx"]
            Zonal wind stress at T-points [N/m^2], ghost ring filled.
        tau_y : Float[Array, "Ny Nx"]
            Meridional wind stress at T-points [N/m^2], ghost ring filled.
        H : float or Float[Array, "Ny Nx"]
            Layer thickness [m]: a scalar or an X-point field.

        Returns
        -------
        Float[Array, "Ny Nx"]
            Forcing at X-points [1/s^2], written on ``[1:-1, 1:-1]`` with a
            zero ghost ring.
        """
        if self.mask is not None:
            tau_x = mask_where(tau_x, self.mask.h)
            tau_y = mask_where(tau_y, self.mask.h)
        # tau_x_on_u[j, i+1/2] = 1/2 (tau_x[j, i] + tau_x[j, i+1]), every row
        tau_x_on_u = (
            jnp.zeros_like(tau_x).at[:, :-1].set(0.5 * (tau_x[:, :-1] + tau_x[:, 1:]))
        )  # [Ny, Nx]; column Nx-1 is outside the domain
        # tau_y_on_v[j+1/2, i] = 1/2 (tau_y[j, i] + tau_y[j+1, i]), every column
        tau_y_on_v = (
            jnp.zeros_like(tau_y).at[:-1, :].set(0.5 * (tau_y[:-1, :] + tau_y[1:, :]))
        )  # [Ny, Nx]; row Ny-1 is outside the domain
        # curl[j+1/2, i+1/2] = d(tau_y)/dx - d(tau_x)/dy  at X-points
        curl = _curl_2d(tau_x_on_u, tau_y_on_v, self.grid.dx, self.grid.dy)
        mx = None if self.mask is None else self.mask.xy_corner_strict
        # H on the X interior; 1 at dry corners so a zero dry thickness cannot give 0/0
        h_int = safe_denominator(interior_or_scalar(H), mx)
        # F[j+1/2, i+1/2] = curl / (rho0 H)
        out = interior(curl[1:-1, 1:-1] / (self.rho0 * h_int), tau_x)
        return mask_where(out, mx)

__call__(tau_x, tau_y, H)

Wind-curl forcing curl(tau) / (rho0 H) at X-points.

Parameters:

Name Type Description Default
tau_x Float[Array, 'Ny Nx']

Zonal wind stress at T-points [N/m^2], ghost ring filled.

 required
tau_y Float[Array, 'Ny Nx']

Meridional wind stress at T-points [N/m^2], ghost ring filled.

 required
H float or Float[Array, 'Ny Nx']

Layer thickness [m]: a scalar or an X-point field.

 required

Returns:

Type Description
Float[Array, 'Ny Nx']

Forcing at X-points [1/s^2], written on [1:-1, 1:-1] with a zero ghost ring.
Source code in finitevolx/_src/forcing/wind.py 
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
def __call__(
    self,
    tau_x: Float[Array, "Ny Nx"],
    tau_y: Float[Array, "Ny Nx"],
    H: float | Float[Array, "Ny Nx"],
) -> Float[Array, "Ny Nx"]:
    """Wind-curl forcing ``curl(tau) / (rho0 H)`` at X-points.

    Parameters
    ----------
    tau_x : Float[Array, "Ny Nx"]
        Zonal wind stress at T-points [N/m^2], ghost ring filled.
    tau_y : Float[Array, "Ny Nx"]
        Meridional wind stress at T-points [N/m^2], ghost ring filled.
    H : float or Float[Array, "Ny Nx"]
        Layer thickness [m]: a scalar or an X-point field.

    Returns
    -------
    Float[Array, "Ny Nx"]
        Forcing at X-points [1/s^2], written on ``[1:-1, 1:-1]`` with a
        zero ghost ring.
    """
    if self.mask is not None:
        tau_x = mask_where(tau_x, self.mask.h)
        tau_y = mask_where(tau_y, self.mask.h)
    # tau_x_on_u[j, i+1/2] = 1/2 (tau_x[j, i] + tau_x[j, i+1]), every row
    tau_x_on_u = (
        jnp.zeros_like(tau_x).at[:, :-1].set(0.5 * (tau_x[:, :-1] + tau_x[:, 1:]))
    )  # [Ny, Nx]; column Nx-1 is outside the domain
    # tau_y_on_v[j+1/2, i] = 1/2 (tau_y[j, i] + tau_y[j+1, i]), every column
    tau_y_on_v = (
        jnp.zeros_like(tau_y).at[:-1, :].set(0.5 * (tau_y[:-1, :] + tau_y[1:, :]))
    )  # [Ny, Nx]; row Ny-1 is outside the domain
    # curl[j+1/2, i+1/2] = d(tau_y)/dx - d(tau_x)/dy  at X-points
    curl = _curl_2d(tau_x_on_u, tau_y_on_v, self.grid.dx, self.grid.dy)
    mx = None if self.mask is None else self.mask.xy_corner_strict
    # H on the X interior; 1 at dry corners so a zero dry thickness cannot give 0/0
    h_int = safe_denominator(interior_or_scalar(H), mx)
    # F[j+1/2, i+1/2] = curl / (rho0 H)
    out = interior(curl[1:-1, 1:-1] / (self.rho0 * h_int), tau_x)
    return mask_where(out, mx)

finitevolx.LinearDrag2D

Bases: AbstractForcing

Linear drag on a 2-D Arakawa C-grid.

du_drag[j, i+1/2] = -r_on_u[j, i+1/2] * u[j, i+1/2]
dv_drag[j+1/2, i] = -r_on_v[j+1/2, i] * v[j+1/2, i]

A spatially varying coefficient r is given at T-points and interpolated to U/V-points with simple x/y averaging.

Parameters:

Name Type Description Default
grid CartesianGrid2D

The underlying 2-D grid.

 required
mask Mask2D or None

Optional land/ocean mask. When provided, du_drag is zeroed at dry U-faces and dv_drag at dry V-faces (via jnp.where, so NaN-filled land never leaks into the output).

 None

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid2D, LinearDrag2D
>>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
>>> drag = LinearDrag2D(grid=grid)
>>> u = jnp.ones((grid.Ny, grid.Nx))
>>> v = jnp.ones((grid.Ny, grid.Nx))
>>> du_drag, dv_drag = drag(u, v, r=1e-7)
Source code in finitevolx/_src/forcing/drag.py 
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
class LinearDrag2D(AbstractForcing):
    """Linear drag on a 2-D Arakawa C-grid.

        du_drag[j, i+1/2] = -r_on_u[j, i+1/2] * u[j, i+1/2]
        dv_drag[j+1/2, i] = -r_on_v[j+1/2, i] * v[j+1/2, i]

    A spatially varying coefficient ``r`` is given at T-points and
    interpolated to U/V-points with simple x/y averaging.

    Parameters
    ----------
    grid : CartesianGrid2D
        The underlying 2-D grid.
    mask : Mask2D or None, optional
        Optional land/ocean mask.  When provided, ``du_drag`` is zeroed at
        dry U-faces and ``dv_drag`` at dry V-faces (via ``jnp.where``, so
        NaN-filled land never leaks into the output).

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid2D, LinearDrag2D
    >>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
    >>> drag = LinearDrag2D(grid=grid)
    >>> u = jnp.ones((grid.Ny, grid.Nx))
    >>> v = jnp.ones((grid.Ny, grid.Nx))
    >>> du_drag, dv_drag = drag(u, v, r=1e-7)
    """

    grid: CartesianGrid2D
    mask: Mask2D | None
    interp: Interpolation2D

    def __init__(
        self,
        grid: CartesianGrid2D,
        mask: Mask2D | None = None,
    ) -> None:
        self.grid = grid
        self.mask = mask
        self.interp = Interpolation2D(grid=grid)

    def __call__(
        self,
        u: Float[Array, "Ny Nx"],
        v: Float[Array, "Ny Nx"],
        r: float | Float[Array, "Ny Nx"],
    ) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
        """Linear drag tendencies (du_drag, dv_drag).

        Parameters
        ----------
        u : Float[Array, "Ny Nx"]
            x-velocity at U-points.
        v : Float[Array, "Ny Nx"]
            y-velocity at V-points.
        r : float or Float[Array, "Ny Nx"]
            Drag coefficient [1/s] — a scalar, or a T-point field (ghost
            ring filled) interpolated to U/V-points.

        Returns
        -------
        tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]
            ``(du_drag, dv_drag)`` at U- and V-points, zero in the ghost
            ring.  When ``self.mask`` is set, dry faces are zeroed.
        """
        if self.mask is not None:
            # Zero dry velocities first: NaN-filled land would otherwise give
            # NaN gradients through -r * u even where the output is masked.
            u = mask_where(u, self.mask.u)
            v = mask_where(v, self.mask.v)

        # r_on_u[j, i+1/2] = 1/2 * (r[j, i] + r[j, i+1])
        r_on_u = on_face_interior(r, self.interp.T_to_U)
        # r_on_v[j+1/2, i] = 1/2 * (r[j, i] + r[j+1, i])
        r_on_v = on_face_interior(r, self.interp.T_to_V)

        # du_drag[j, i+1/2] = -r_on_u * u
        du_drag = interior(linear_drag_tendency(u[1:-1, 1:-1], r_on_u), u)
        # dv_drag[j+1/2, i] = -r_on_v * v
        dv_drag = interior(linear_drag_tendency(v[1:-1, 1:-1], r_on_v), v)

        if self.mask is None:
            return du_drag, dv_drag
        return mask_where(du_drag, self.mask.u), mask_where(dv_drag, self.mask.v)

__call__(u, v, r)

Linear drag tendencies (du_drag, dv_drag).

Parameters:

Name Type Description Default
u Float[Array, 'Ny Nx']

x-velocity at U-points.

 required
v Float[Array, 'Ny Nx']

y-velocity at V-points.

 required
r float or Float[Array, 'Ny Nx']

Drag coefficient [1/s] — a scalar, or a T-point field (ghost ring filled) interpolated to U/V-points.

 required

Returns:

Type Description
tuple[Float[Array, 'Ny Nx'], Float[Array, 'Ny Nx']]

(du_drag, dv_drag) at U- and V-points, zero in the ghost ring. When self.mask is set, dry faces are zeroed.
Source code in finitevolx/_src/forcing/drag.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def __call__(
    self,
    u: Float[Array, "Ny Nx"],
    v: Float[Array, "Ny Nx"],
    r: float | Float[Array, "Ny Nx"],
) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
    """Linear drag tendencies (du_drag, dv_drag).

    Parameters
    ----------
    u : Float[Array, "Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "Ny Nx"]
        y-velocity at V-points.
    r : float or Float[Array, "Ny Nx"]
        Drag coefficient [1/s] — a scalar, or a T-point field (ghost
        ring filled) interpolated to U/V-points.

    Returns
    -------
    tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]
        ``(du_drag, dv_drag)`` at U- and V-points, zero in the ghost
        ring.  When ``self.mask`` is set, dry faces are zeroed.
    """
    if self.mask is not None:
        # Zero dry velocities first: NaN-filled land would otherwise give
        # NaN gradients through -r * u even where the output is masked.
        u = mask_where(u, self.mask.u)
        v = mask_where(v, self.mask.v)

    # r_on_u[j, i+1/2] = 1/2 * (r[j, i] + r[j, i+1])
    r_on_u = on_face_interior(r, self.interp.T_to_U)
    # r_on_v[j+1/2, i] = 1/2 * (r[j, i] + r[j+1, i])
    r_on_v = on_face_interior(r, self.interp.T_to_V)

    # du_drag[j, i+1/2] = -r_on_u * u
    du_drag = interior(linear_drag_tendency(u[1:-1, 1:-1], r_on_u), u)
    # dv_drag[j+1/2, i] = -r_on_v * v
    dv_drag = interior(linear_drag_tendency(v[1:-1, 1:-1], r_on_v), v)

    if self.mask is None:
        return du_drag, dv_drag
    return mask_where(du_drag, self.mask.u), mask_where(dv_drag, self.mask.v)

finitevolx.QuadraticDrag2D

Bases: AbstractForcing

Quadratic bottom drag on a 2-D Arakawa C-grid.

du_drag[j, i+1/2] = -Cd * |u|_on_u * u[j, i+1/2] / h_bot_on_u
dv_drag[j+1/2, i] = -Cd * |u|_on_v * v[j+1/2, i] / h_bot_on_v

The flow speed is evaluated at each face from the native component and the 4-point average of the cross component:

|u|_on_u[j, i+1/2] = sqrt(u[j, i+1/2]^2 + v_on_u[j, i+1/2]^2)
|u|_on_v[j+1/2, i] = sqrt(v[j+1/2, i]^2 + u_on_v[j+1/2, i]^2)

so boundary faces see the caller-filled ghost velocities rather than a zero ghost-ring speed. The speed is gradient-safe at rest. Spatially varying cd / h_bot are given at T-points and interpolated to U/V-points with simple x/y averaging.

Parameters:

Name Type Description Default
grid CartesianGrid2D

The underlying 2-D grid.

 required
mask Mask2D or None

Optional land/ocean mask. When provided, du_drag is zeroed at dry U-faces and dv_drag at dry V-faces (via jnp.where, so NaN-filled land never leaks into the output).

 None

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid2D, QuadraticDrag2D
>>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
>>> drag = QuadraticDrag2D(grid=grid)
>>> u = jnp.ones((grid.Ny, grid.Nx))
>>> v = jnp.zeros((grid.Ny, grid.Nx))
>>> du_drag, dv_drag = drag(u, v, cd=2.5e-3, h_bot=100.0)
Source code in finitevolx/_src/forcing/drag.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
class QuadraticDrag2D(AbstractForcing):
    """Quadratic bottom drag on a 2-D Arakawa C-grid.

        du_drag[j, i+1/2] = -Cd * |u|_on_u * u[j, i+1/2] / h_bot_on_u
        dv_drag[j+1/2, i] = -Cd * |u|_on_v * v[j+1/2, i] / h_bot_on_v

    The flow speed is evaluated **at each face** from the native component
    and the 4-point average of the cross component:

        |u|_on_u[j, i+1/2] = sqrt(u[j, i+1/2]^2 + v_on_u[j, i+1/2]^2)
        |u|_on_v[j+1/2, i] = sqrt(v[j+1/2, i]^2 + u_on_v[j+1/2, i]^2)

    so boundary faces see the caller-filled ghost velocities rather than a
    zero ghost-ring speed.  The speed is gradient-safe at rest.  Spatially
    varying ``cd`` / ``h_bot`` are given at T-points and interpolated to
    U/V-points with simple x/y averaging.

    Parameters
    ----------
    grid : CartesianGrid2D
        The underlying 2-D grid.
    mask : Mask2D or None, optional
        Optional land/ocean mask.  When provided, ``du_drag`` is zeroed at
        dry U-faces and ``dv_drag`` at dry V-faces (via ``jnp.where``, so
        NaN-filled land never leaks into the output).

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid2D, QuadraticDrag2D
    >>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
    >>> drag = QuadraticDrag2D(grid=grid)
    >>> u = jnp.ones((grid.Ny, grid.Nx))
    >>> v = jnp.zeros((grid.Ny, grid.Nx))
    >>> du_drag, dv_drag = drag(u, v, cd=2.5e-3, h_bot=100.0)
    """

    grid: CartesianGrid2D
    mask: Mask2D | None
    interp: Interpolation2D

    def __init__(
        self,
        grid: CartesianGrid2D,
        mask: Mask2D | None = None,
    ) -> None:
        self.grid = grid
        self.mask = mask
        self.interp = Interpolation2D(grid=grid)

    def __call__(
        self,
        u: Float[Array, "Ny Nx"],
        v: Float[Array, "Ny Nx"],
        cd: float | Float[Array, "Ny Nx"],
        h_bot: float | Float[Array, "Ny Nx"],
    ) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
        """Quadratic drag tendencies (du_drag, dv_drag).

        Parameters
        ----------
        u : Float[Array, "Ny Nx"]
            x-velocity at U-points.
        v : Float[Array, "Ny Nx"]
            y-velocity at V-points.
        cd : float or Float[Array, "Ny Nx"]
            Dimensionless drag coefficient — a scalar, or a T-point field
            (ghost ring filled) interpolated to U/V-points.
        h_bot : float or Float[Array, "Ny Nx"]
            Bottom-layer thickness [m] — a scalar, or a T-point field
            (ghost ring filled) interpolated to U/V-points.

        Returns
        -------
        tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]
            ``(du_drag, dv_drag)`` at U- and V-points, zero in the ghost
            ring.  When ``self.mask`` is set, dry faces are zeroed.
        """
        mu = None if self.mask is None else self.mask.u
        mv = None if self.mask is None else self.mask.v
        return self._tendency(u, v, cd, h_bot, mu, mv)

    def _tendency(
        self,
        u: Float[Array, "Ny Nx"],
        v: Float[Array, "Ny Nx"],
        cd: float | Float[Array, "Ny Nx"],
        h_bot: float | Float[Array, "Ny Nx"],
        mu: Array | None,
        mv: Array | None,
    ) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
        """Core stencil with explicit ``[Ny, Nx]`` face masks (or ``None``).

        Shared with :class:`QuadraticDrag3D`, which passes one z-level's masks.
        """
        # Zero dry velocities first (land velocity is zero): the 4-point
        # cross-face average at a wet coastal face reads dry neighbours, and
        # NaN-filled land would otherwise make the speed there NaN.
        u = mask_where(u, mu)
        v = mask_where(v, mv)

        # v_on_u[j, i+1/2] = 1/4*(v[j+1/2,i] + v[j-1/2,i] + v[j+1/2,i+1] + v[j-1/2,i+1])
        v_on_u = self.interp.V_to_U(v)[1:-1, 1:-1]
        # u_on_v[j+1/2, i] = 1/4*(u[j,i+1/2] + u[j+1,i+1/2] + u[j,i-1/2] + u[j+1,i-1/2])
        u_on_v = self.interp.U_to_V(u)[1:-1, 1:-1]
        u_in = u[1:-1, 1:-1]
        v_in = v[1:-1, 1:-1]

        # |u|_on_u[j, i+1/2] = sqrt(u^2 + v_on_u^2)
        speed_u = safe_speed(u_in, v_on_u)
        # |u|_on_v[j+1/2, i] = sqrt(v^2 + u_on_v^2)
        speed_v = safe_speed(v_in, u_on_v)

        cd_on_u = on_face_interior(cd, self.interp.T_to_U)
        cd_on_v = on_face_interior(cd, self.interp.T_to_V)
        # Swap in 1 at dry faces so a zero dry-cell thickness cannot give 0/0.
        h_on_u = safe_denominator(on_face_interior(h_bot, self.interp.T_to_U), mu)
        h_on_v = safe_denominator(on_face_interior(h_bot, self.interp.T_to_V), mv)

        # du_drag[j, i+1/2] = -Cd * |u|_on_u * u / h_bot_on_u
        du_drag = interior(quadratic_drag_tendency(u_in, speed_u, cd_on_u, h_on_u), u)
        # dv_drag[j+1/2, i] = -Cd * |u|_on_v * v / h_bot_on_v
        dv_drag = interior(quadratic_drag_tendency(v_in, speed_v, cd_on_v, h_on_v), v)

        return mask_where(du_drag, mu), mask_where(dv_drag, mv)

__call__(u, v, cd, h_bot)

Quadratic drag tendencies (du_drag, dv_drag).

Parameters:

Name Type Description Default
u Float[Array, 'Ny Nx']

x-velocity at U-points.

 required
v Float[Array, 'Ny Nx']

y-velocity at V-points.

 required
cd float or Float[Array, 'Ny Nx']

Dimensionless drag coefficient — a scalar, or a T-point field (ghost ring filled) interpolated to U/V-points.

 required
h_bot float or Float[Array, 'Ny Nx']

Bottom-layer thickness [m] — a scalar, or a T-point field (ghost ring filled) interpolated to U/V-points.

 required

Returns:

Type Description
tuple[Float[Array, 'Ny Nx'], Float[Array, 'Ny Nx']]

(du_drag, dv_drag) at U- and V-points, zero in the ghost ring. When self.mask is set, dry faces are zeroed.
Source code in finitevolx/_src/forcing/drag.py 
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
def __call__(
    self,
    u: Float[Array, "Ny Nx"],
    v: Float[Array, "Ny Nx"],
    cd: float | Float[Array, "Ny Nx"],
    h_bot: float | Float[Array, "Ny Nx"],
) -> tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]:
    """Quadratic drag tendencies (du_drag, dv_drag).

    Parameters
    ----------
    u : Float[Array, "Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "Ny Nx"]
        y-velocity at V-points.
    cd : float or Float[Array, "Ny Nx"]
        Dimensionless drag coefficient — a scalar, or a T-point field
        (ghost ring filled) interpolated to U/V-points.
    h_bot : float or Float[Array, "Ny Nx"]
        Bottom-layer thickness [m] — a scalar, or a T-point field
        (ghost ring filled) interpolated to U/V-points.

    Returns
    -------
    tuple[Float[Array, "Ny Nx"], Float[Array, "Ny Nx"]]
        ``(du_drag, dv_drag)`` at U- and V-points, zero in the ghost
        ring.  When ``self.mask`` is set, dry faces are zeroed.
    """
    mu = None if self.mask is None else self.mask.u
    mv = None if self.mask is None else self.mask.v
    return self._tendency(u, v, cd, h_bot, mu, mv)

finitevolx.RayleighDamping2D

Bases: AbstractForcing

Rayleigh damping of a T-point field on a 2-D Arakawa C-grid.

dq[j, i] = -r[j, i] * (q[j, i] - q_ref[j, i])

Works on any scalar field at T-points (layer thickness, tracer, ...). With q_ref=None this is pure linear damping toward zero.

Parameters:

Name Type Description Default
grid CartesianGrid2D

The underlying 2-D grid.

 required
mask Mask2D or None

Optional land/ocean mask. When provided, dq is zeroed at dry T-cells (via jnp.where, so NaN-filled land never leaks into the output).

 required

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid2D, RayleighDamping2D
>>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
>>> damp = RayleighDamping2D(grid=grid)
>>> h = jnp.ones((grid.Ny, grid.Nx))
>>> dh = damp(h, r=1e-5, q_ref=0.5 * h)
Source code in finitevolx/_src/forcing/rayleigh.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
class RayleighDamping2D(AbstractForcing):
    """Rayleigh damping of a T-point field on a 2-D Arakawa C-grid.

        dq[j, i] = -r[j, i] * (q[j, i] - q_ref[j, i])

    Works on any scalar field at T-points (layer thickness, tracer, ...).
    With ``q_ref=None`` this is pure linear damping toward zero.

    Parameters
    ----------
    grid : CartesianGrid2D
        The underlying 2-D grid.
    mask : Mask2D or None, optional
        Optional land/ocean mask.  When provided, ``dq`` is zeroed at dry
        T-cells (via ``jnp.where``, so NaN-filled land never leaks into the
        output).

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid2D, RayleighDamping2D
    >>> grid = CartesianGrid2D.from_interior(8, 8, 1.0, 1.0)
    >>> damp = RayleighDamping2D(grid=grid)
    >>> h = jnp.ones((grid.Ny, grid.Nx))
    >>> dh = damp(h, r=1e-5, q_ref=0.5 * h)
    """

    grid: CartesianGrid2D
    mask: Mask2D | None = None

    def __call__(
        self,
        q: Float[Array, "Ny Nx"],
        r: float | Float[Array, "Ny Nx"],
        q_ref: float | Float[Array, "Ny Nx"] | None = None,
    ) -> Float[Array, "Ny Nx"]:
        """Rayleigh damping tendency at T-points.

        Parameters
        ----------
        q : Float[Array, "Ny Nx"]
            Field at T-points.
        r : float or Float[Array, "Ny Nx"]
            Damping rate [1/s] — a scalar or a T-point field.
        q_ref : float, Float[Array, "Ny Nx"] or None, optional
            Reference state — a scalar or a T-point field.  ``None``
            (default) damps toward zero.

        Returns
        -------
        Float[Array, "Ny Nx"]
            Damping tendency at T-points, zero in the ghost ring.  When
            ``self.mask`` is set, dry cells are zeroed.
        """
        if self.mask is not None:
            # Zero dry cells first: NaN-filled land would otherwise give NaN
            # gradients through -r * (q - q_ref) even where dq is masked.
            q = mask_where(q, self.mask.h)
            if q_ref is not None and jnp.ndim(q_ref) == 2:
                q_ref = mask_where(jnp.asarray(q_ref), self.mask.h)

        q_ref_in = None if q_ref is None else interior_or_scalar(q_ref)
        # dq[j, i] = -r[j, i] * (q[j, i] - q_ref[j, i])
        dq = interior(
            rayleigh_tendency(q[1:-1, 1:-1], interior_or_scalar(r), q_ref_in), q
        )

        if self.mask is None:
            return dq
        return mask_where(dq, self.mask.h)

__call__(q, r, q_ref=None)

Rayleigh damping tendency at T-points.

Parameters:

Name Type Description Default
q Float[Array, 'Ny Nx']

Field at T-points.

 required
r float or Float[Array, 'Ny Nx']

Damping rate [1/s] — a scalar or a T-point field.

 required
q_ref (float, Float[Array, 'Ny Nx'] or None)

Reference state — a scalar or a T-point field. None (default) damps toward zero.

 None

Returns:

Type Description
Float[Array, 'Ny Nx']

Damping tendency at T-points, zero in the ghost ring. When self.mask is set, dry cells are zeroed.
Source code in finitevolx/_src/forcing/rayleigh.py 
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def __call__(
    self,
    q: Float[Array, "Ny Nx"],
    r: float | Float[Array, "Ny Nx"],
    q_ref: float | Float[Array, "Ny Nx"] | None = None,
) -> Float[Array, "Ny Nx"]:
    """Rayleigh damping tendency at T-points.

    Parameters
    ----------
    q : Float[Array, "Ny Nx"]
        Field at T-points.
    r : float or Float[Array, "Ny Nx"]
        Damping rate [1/s] — a scalar or a T-point field.
    q_ref : float, Float[Array, "Ny Nx"] or None, optional
        Reference state — a scalar or a T-point field.  ``None``
        (default) damps toward zero.

    Returns
    -------
    Float[Array, "Ny Nx"]
        Damping tendency at T-points, zero in the ghost ring.  When
        ``self.mask`` is set, dry cells are zeroed.
    """
    if self.mask is not None:
        # Zero dry cells first: NaN-filled land would otherwise give NaN
        # gradients through -r * (q - q_ref) even where dq is masked.
        q = mask_where(q, self.mask.h)
        if q_ref is not None and jnp.ndim(q_ref) == 2:
            q_ref = mask_where(jnp.asarray(q_ref), self.mask.h)

    q_ref_in = None if q_ref is None else interior_or_scalar(q_ref)
    # dq[j, i] = -r[j, i] * (q[j, i] - q_ref[j, i])
    dq = interior(
        rayleigh_tendency(q[1:-1, 1:-1], interior_or_scalar(r), q_ref_in), q
    )

    if self.mask is None:
        return dq
    return mask_where(dq, self.mask.h)

3D operators

Wrappers for CartesianGrid3D built on the 2D operators, with an optional Mask3D. Levels are stacked bottom-up — k = 0 is the bottom ghost and k = Nz - 1 the top ghost. Each operator follows one of three vertical patterns:

Pattern Operators Levels written
Top injection WindStress3D top interior level k = Nz - 2 only
Bottom injection LinearDrag3D, QuadraticDrag3D bottom interior level k = 1 only
Uniform vmap RayleighDamping3D every interior level (z-ghosts zeroed)

finitevolx.WindStress3D

Bases: AbstractForcing

Surface wind-stress forcing on a 3-D Arakawa C-grid.

Applies the :class:WindStress2D stencil to the top interior z-level only and returns zero at every other level. The 3-D grid stacks levels bottom-up (k = 0 bottom ghost, k = Nz - 1 top ghost), so the top interior level is k_top = Nz - 2:

du_wind[k_top, j, i+1/2] = tau_x_on_u[j, i+1/2] / (rho0 * dz_top_on_u[j, i+1/2])
dv_wind[k_top, j+1/2, i] = tau_y_on_v[j+1/2, i] / (rho0 * dz_top_on_v[j+1/2, i])
du_wind[k != k_top] = dv_wind[k != k_top] = 0

Parameters:

Name Type Description Default
grid CartesianGrid3D

The underlying 3-D grid.

 required
mask Mask3D or None

Optional land/ocean mask. The k_top slices of mask.u / mask.v zero dry faces (via jnp.where) and guard the dz_top division there, so zero dry-cell thickness or NaN-filled land never leaks NaN.

 None
rho0 float

Reference density [kg/m^3]. Default 1025.0.

 1025.0

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid3D, WindStress3D
>>> grid = CartesianGrid3D.from_interior(6, 6, 4, 1.0, 1.0, 1.0)
>>> wind = WindStress3D(grid=grid)
>>> tau = 0.1 * jnp.ones((grid.Ny, grid.Nx))
>>> du_wind, dv_wind = wind(tau, tau, dz_top=grid.dz)
>>> du_wind.shape
(6, 8, 8)
Source code in finitevolx/_src/forcing/wind.py 
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
class WindStress3D(AbstractForcing):
    """Surface wind-stress forcing on a 3-D Arakawa C-grid.

    Applies the :class:`WindStress2D` stencil to the **top interior
    z-level only** and returns zero at every other level.  The 3-D grid
    stacks levels bottom-up (``k = 0`` bottom ghost, ``k = Nz - 1`` top
    ghost), so the top interior level is ``k_top = Nz - 2``:

        du_wind[k_top, j, i+1/2] = tau_x_on_u[j, i+1/2] / (rho0 * dz_top_on_u[j, i+1/2])
        dv_wind[k_top, j+1/2, i] = tau_y_on_v[j+1/2, i] / (rho0 * dz_top_on_v[j+1/2, i])
        du_wind[k != k_top] = dv_wind[k != k_top] = 0

    Parameters
    ----------
    grid : CartesianGrid3D
        The underlying 3-D grid.
    mask : Mask3D or None, optional
        Optional land/ocean mask.  The ``k_top`` slices of ``mask.u`` /
        ``mask.v`` zero dry faces (via ``jnp.where``) and guard the
        ``dz_top`` division there, so zero dry-cell thickness or NaN-filled
        land never leaks NaN.
    rho0 : float, optional
        Reference density [kg/m^3].  Default ``1025.0``.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid3D, WindStress3D
    >>> grid = CartesianGrid3D.from_interior(6, 6, 4, 1.0, 1.0, 1.0)
    >>> wind = WindStress3D(grid=grid)
    >>> tau = 0.1 * jnp.ones((grid.Ny, grid.Nx))
    >>> du_wind, dv_wind = wind(tau, tau, dz_top=grid.dz)
    >>> du_wind.shape
    (6, 8, 8)
    """

    grid: CartesianGrid3D
    mask: Mask3D | None
    _wind2d: WindStress2D

    def __init__(
        self,
        grid: CartesianGrid3D,
        mask: Mask3D | None = None,
        rho0: float = 1025.0,
    ) -> None:
        self.grid = grid
        self.mask = mask
        # The inner 2-D op is mask-free; the 3-D wrapper passes per-level masks.
        self._wind2d = WindStress2D(grid=grid.horizontal_grid(), rho0=rho0)

    def __call__(
        self,
        tau_x: Float[Array, "Ny Nx"],
        tau_y: Float[Array, "Ny Nx"],
        dz_top: float | Float[Array, "Ny Nx"],
    ) -> tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]:
        """Wind-stress tendencies, non-zero only at the top interior level.

        Parameters
        ----------
        tau_x : Float[Array, "Ny Nx"]
            Zonal wind stress at T-points [N/m^2].
        tau_y : Float[Array, "Ny Nx"]
            Meridional wind stress at T-points [N/m^2].
        dz_top : float or Float[Array, "Ny Nx"]
            Top-layer thickness [m] — a scalar, or a T-point field (ghost
            ring filled) interpolated to U/V-points.

        Returns
        -------
        tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]
            ``(du_wind, dv_wind)`` at U- and V-points, zero outside
            ``k = Nz - 2`` and in the ghost ring.  When ``self.mask`` is set,
            dry faces are zeroed.
        """
        k_top = self.grid.Nz - 2
        mu = None if self.mask is None else self.mask.u[k_top]
        mv = None if self.mask is None else self.mask.v[k_top]
        du_2d, dv_2d = self._wind2d._tendency(tau_x, tau_y, dz_top, mu, mv)

        # Inject at the top interior z-level k_top = Nz - 2.
        shape = (self.grid.Nz, self.grid.Ny, self.grid.Nx)
        du_wind = jnp.zeros(shape, dtype=du_2d.dtype).at[k_top].set(du_2d)
        dv_wind = jnp.zeros(shape, dtype=dv_2d.dtype).at[k_top].set(dv_2d)
        return du_wind, dv_wind

__call__(tau_x, tau_y, dz_top)

Wind-stress tendencies, non-zero only at the top interior level.

Parameters:

Name Type Description Default
tau_x Float[Array, 'Ny Nx']

Zonal wind stress at T-points [N/m^2].

 required
tau_y Float[Array, 'Ny Nx']

Meridional wind stress at T-points [N/m^2].

 required
dz_top float or Float[Array, 'Ny Nx']

Top-layer thickness [m] — a scalar, or a T-point field (ghost ring filled) interpolated to U/V-points.

 required

Returns:

Type Description
tuple[Float[Array, 'Nz Ny Nx'], Float[Array, 'Nz Ny Nx']]

(du_wind, dv_wind) at U- and V-points, zero outside k = Nz - 2 and in the ghost ring. When self.mask is set, dry faces are zeroed.
Source code in finitevolx/_src/forcing/wind.py 
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
def __call__(
    self,
    tau_x: Float[Array, "Ny Nx"],
    tau_y: Float[Array, "Ny Nx"],
    dz_top: float | Float[Array, "Ny Nx"],
) -> tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]:
    """Wind-stress tendencies, non-zero only at the top interior level.

    Parameters
    ----------
    tau_x : Float[Array, "Ny Nx"]
        Zonal wind stress at T-points [N/m^2].
    tau_y : Float[Array, "Ny Nx"]
        Meridional wind stress at T-points [N/m^2].
    dz_top : float or Float[Array, "Ny Nx"]
        Top-layer thickness [m] — a scalar, or a T-point field (ghost
        ring filled) interpolated to U/V-points.

    Returns
    -------
    tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]
        ``(du_wind, dv_wind)`` at U- and V-points, zero outside
        ``k = Nz - 2`` and in the ghost ring.  When ``self.mask`` is set,
        dry faces are zeroed.
    """
    k_top = self.grid.Nz - 2
    mu = None if self.mask is None else self.mask.u[k_top]
    mv = None if self.mask is None else self.mask.v[k_top]
    du_2d, dv_2d = self._wind2d._tendency(tau_x, tau_y, dz_top, mu, mv)

    # Inject at the top interior z-level k_top = Nz - 2.
    shape = (self.grid.Nz, self.grid.Ny, self.grid.Nx)
    du_wind = jnp.zeros(shape, dtype=du_2d.dtype).at[k_top].set(du_2d)
    dv_wind = jnp.zeros(shape, dtype=dv_2d.dtype).at[k_top].set(dv_2d)
    return du_wind, dv_wind

finitevolx.LinearDrag3D

Bases: AbstractForcing

Linear bottom drag on a 3-D Arakawa C-grid.

Applies the :class:LinearDrag2D stencil to the bottom interior z-level only and returns zero at every other level. The 3-D grid stacks levels bottom-up (k = 0 bottom ghost, k = Nz - 1 top ghost), so the bottom interior level is k_bot = 1:

du_drag[1, j, i+1/2] = -r_on_u[j, i+1/2] * u[1, j, i+1/2]
dv_drag[1, j+1/2, i] = -r_on_v[j+1/2, i] * v[1, j+1/2, i]

The bottom is the lowest interior level everywhere; a per-column sea-floor index (e.g. Mask3D.k_bottom) is not used.

Parameters:

Name Type Description Default
grid CartesianGrid3D

The underlying 3-D grid.

 required
mask Mask3D or None

Optional land/ocean mask. The k_bot slices of mask.u / mask.v zero dry faces via jnp.where, so NaN-filled land never leaks into the output.

 None

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid3D, LinearDrag3D
>>> grid = CartesianGrid3D.from_interior(6, 6, 4, 1.0, 1.0, 1.0)
>>> drag = LinearDrag3D(grid=grid)
>>> u = jnp.ones((grid.Nz, grid.Ny, grid.Nx))
>>> du_drag, dv_drag = drag(u, u, r=1e-7)
Source code in finitevolx/_src/forcing/drag.py 
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
class LinearDrag3D(AbstractForcing):
    """Linear bottom drag on a 3-D Arakawa C-grid.

    Applies the :class:`LinearDrag2D` stencil to the **bottom interior
    z-level only** and returns zero at every other level.  The 3-D grid
    stacks levels bottom-up (``k = 0`` bottom ghost, ``k = Nz - 1`` top
    ghost), so the bottom interior level is ``k_bot = 1``:

        du_drag[1, j, i+1/2] = -r_on_u[j, i+1/2] * u[1, j, i+1/2]
        dv_drag[1, j+1/2, i] = -r_on_v[j+1/2, i] * v[1, j+1/2, i]

    The bottom is the lowest interior level everywhere; a per-column
    sea-floor index (e.g. ``Mask3D.k_bottom``) is not used.

    Parameters
    ----------
    grid : CartesianGrid3D
        The underlying 3-D grid.
    mask : Mask3D or None, optional
        Optional land/ocean mask.  The ``k_bot`` slices of ``mask.u`` /
        ``mask.v`` zero dry faces via ``jnp.where``, so NaN-filled land
        never leaks into the output.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid3D, LinearDrag3D
    >>> grid = CartesianGrid3D.from_interior(6, 6, 4, 1.0, 1.0, 1.0)
    >>> drag = LinearDrag3D(grid=grid)
    >>> u = jnp.ones((grid.Nz, grid.Ny, grid.Nx))
    >>> du_drag, dv_drag = drag(u, u, r=1e-7)
    """

    grid: CartesianGrid3D
    mask: Mask3D | None
    _drag2d: LinearDrag2D

    def __init__(
        self,
        grid: CartesianGrid3D,
        mask: Mask3D | None = None,
    ) -> None:
        self.grid = grid
        self.mask = mask
        # The inner 2-D op is mask-free; the 3-D wrapper passes per-level masks.
        self._drag2d = LinearDrag2D(grid=grid.horizontal_grid())

    def __call__(
        self,
        u: Float[Array, "Nz Ny Nx"],
        v: Float[Array, "Nz Ny Nx"],
        r: float | Float[Array, "Ny Nx"],
    ) -> tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]:
        """Linear drag tendencies, non-zero only at the bottom interior level.

        Parameters
        ----------
        u : Float[Array, "Nz Ny Nx"]
            x-velocity at U-points.
        v : Float[Array, "Nz Ny Nx"]
            y-velocity at V-points.
        r : float or Float[Array, "Ny Nx"]
            Drag coefficient [1/s] — a scalar, or a T-point field (ghost
            ring filled) interpolated to U/V-points.

        Returns
        -------
        tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]
            ``(du_drag, dv_drag)`` at U- and V-points, zero outside
            ``k = 1`` and in the ghost ring.  When ``self.mask`` is
            set, dry faces are zeroed.
        """
        u_k, v_k = u[_K_BOT], v[_K_BOT]
        if self.mask is not None:
            # Zero dry velocities first so NaN-filled land cannot give NaN
            # gradients through -r * u.
            u_k = mask_where(u_k, self.mask.u[_K_BOT])
            v_k = mask_where(v_k, self.mask.v[_K_BOT])
        du_2d, dv_2d = self._drag2d(u_k, v_k, r)
        if self.mask is not None:
            du_2d = mask_where(du_2d, self.mask.u[_K_BOT])
            dv_2d = mask_where(dv_2d, self.mask.v[_K_BOT])
        return _inject_bottom(du_2d, dv_2d, u, v)

__call__(u, v, r)

Linear drag tendencies, non-zero only at the bottom interior level.

Parameters:

Name Type Description Default
u Float[Array, 'Nz Ny Nx']

x-velocity at U-points.

 required
v Float[Array, 'Nz Ny Nx']

y-velocity at V-points.

 required
r float or Float[Array, 'Ny Nx']

Drag coefficient [1/s] — a scalar, or a T-point field (ghost ring filled) interpolated to U/V-points.

 required

Returns:

Type Description
tuple[Float[Array, 'Nz Ny Nx'], Float[Array, 'Nz Ny Nx']]

(du_drag, dv_drag) at U- and V-points, zero outside k = 1 and in the ghost ring. When self.mask is set, dry faces are zeroed.
Source code in finitevolx/_src/forcing/drag.py 
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
def __call__(
    self,
    u: Float[Array, "Nz Ny Nx"],
    v: Float[Array, "Nz Ny Nx"],
    r: float | Float[Array, "Ny Nx"],
) -> tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]:
    """Linear drag tendencies, non-zero only at the bottom interior level.

    Parameters
    ----------
    u : Float[Array, "Nz Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "Nz Ny Nx"]
        y-velocity at V-points.
    r : float or Float[Array, "Ny Nx"]
        Drag coefficient [1/s] — a scalar, or a T-point field (ghost
        ring filled) interpolated to U/V-points.

    Returns
    -------
    tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]
        ``(du_drag, dv_drag)`` at U- and V-points, zero outside
        ``k = 1`` and in the ghost ring.  When ``self.mask`` is
        set, dry faces are zeroed.
    """
    u_k, v_k = u[_K_BOT], v[_K_BOT]
    if self.mask is not None:
        # Zero dry velocities first so NaN-filled land cannot give NaN
        # gradients through -r * u.
        u_k = mask_where(u_k, self.mask.u[_K_BOT])
        v_k = mask_where(v_k, self.mask.v[_K_BOT])
    du_2d, dv_2d = self._drag2d(u_k, v_k, r)
    if self.mask is not None:
        du_2d = mask_where(du_2d, self.mask.u[_K_BOT])
        dv_2d = mask_where(dv_2d, self.mask.v[_K_BOT])
    return _inject_bottom(du_2d, dv_2d, u, v)

finitevolx.QuadraticDrag3D

Bases: AbstractForcing

Quadratic bottom drag on a 3-D Arakawa C-grid.

Applies the :class:QuadraticDrag2D stencil to the bottom interior z-level only and returns zero at every other level. The 3-D grid stacks levels bottom-up (k = 0 bottom ghost, k = Nz - 1 top ghost), so the bottom interior level is k_bot = 1:

du_drag[1, j, i+1/2] = -Cd * |u|_on_u * u[1, j, i+1/2] / h_bot_on_u
dv_drag[1, j+1/2, i] = -Cd * |u|_on_v * v[1, j+1/2, i] / h_bot_on_v

The bottom is the lowest interior level everywhere; a per-column sea-floor index (e.g. Mask3D.k_bottom) is not used.

Parameters:

Name Type Description Default
grid CartesianGrid3D

The underlying 3-D grid.

 required
mask Mask3D or None

Optional land/ocean mask. The k_bot slices of mask.u / mask.v zero dry faces via jnp.where, so NaN-filled land never leaks into the output.

 None

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid3D, QuadraticDrag3D
>>> grid = CartesianGrid3D.from_interior(6, 6, 4, 1.0, 1.0, 1.0)
>>> drag = QuadraticDrag3D(grid=grid)
>>> u = jnp.ones((grid.Nz, grid.Ny, grid.Nx))
>>> du_drag, dv_drag = drag(u, u, cd=2.5e-3, h_bot=grid.dz)
Source code in finitevolx/_src/forcing/drag.py 
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
class QuadraticDrag3D(AbstractForcing):
    """Quadratic bottom drag on a 3-D Arakawa C-grid.

    Applies the :class:`QuadraticDrag2D` stencil to the **bottom interior
    z-level only** and returns zero at every other level.  The 3-D grid
    stacks levels bottom-up (``k = 0`` bottom ghost, ``k = Nz - 1`` top
    ghost), so the bottom interior level is ``k_bot = 1``:

        du_drag[1, j, i+1/2] = -Cd * |u|_on_u * u[1, j, i+1/2] / h_bot_on_u
        dv_drag[1, j+1/2, i] = -Cd * |u|_on_v * v[1, j+1/2, i] / h_bot_on_v

    The bottom is the lowest interior level everywhere; a per-column
    sea-floor index (e.g. ``Mask3D.k_bottom``) is not used.

    Parameters
    ----------
    grid : CartesianGrid3D
        The underlying 3-D grid.
    mask : Mask3D or None, optional
        Optional land/ocean mask.  The ``k_bot`` slices of ``mask.u`` /
        ``mask.v`` zero dry faces via ``jnp.where``, so NaN-filled land
        never leaks into the output.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid3D, QuadraticDrag3D
    >>> grid = CartesianGrid3D.from_interior(6, 6, 4, 1.0, 1.0, 1.0)
    >>> drag = QuadraticDrag3D(grid=grid)
    >>> u = jnp.ones((grid.Nz, grid.Ny, grid.Nx))
    >>> du_drag, dv_drag = drag(u, u, cd=2.5e-3, h_bot=grid.dz)
    """

    grid: CartesianGrid3D
    mask: Mask3D | None
    _drag2d: QuadraticDrag2D

    def __init__(
        self,
        grid: CartesianGrid3D,
        mask: Mask3D | None = None,
    ) -> None:
        self.grid = grid
        self.mask = mask
        # The inner 2-D op is mask-free; the 3-D wrapper passes per-level masks.
        self._drag2d = QuadraticDrag2D(grid=grid.horizontal_grid())

    def __call__(
        self,
        u: Float[Array, "Nz Ny Nx"],
        v: Float[Array, "Nz Ny Nx"],
        cd: float | Float[Array, "Ny Nx"],
        h_bot: float | Float[Array, "Ny Nx"],
    ) -> tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]:
        """Quadratic drag tendencies, non-zero only at the bottom interior level.

        Parameters
        ----------
        u : Float[Array, "Nz Ny Nx"]
            x-velocity at U-points.
        v : Float[Array, "Nz Ny Nx"]
            y-velocity at V-points.
        cd : float or Float[Array, "Ny Nx"]
            Dimensionless drag coefficient — a scalar, or a T-point field
            (ghost ring filled) interpolated to U/V-points.
        h_bot : float or Float[Array, "Ny Nx"]
            Bottom-layer thickness [m] — a scalar, or a T-point field
            (ghost ring filled) interpolated to U/V-points.

        Returns
        -------
        tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]
            ``(du_drag, dv_drag)`` at U- and V-points, zero outside
            ``k = 1`` and in the ghost ring.  When ``self.mask`` is
            set, dry faces are zeroed.
        """
        mu = None if self.mask is None else self.mask.u[_K_BOT]
        mv = None if self.mask is None else self.mask.v[_K_BOT]
        du_2d, dv_2d = self._drag2d._tendency(u[_K_BOT], v[_K_BOT], cd, h_bot, mu, mv)
        return _inject_bottom(du_2d, dv_2d, u, v)

__call__(u, v, cd, h_bot)

Quadratic drag tendencies, non-zero only at the bottom interior level.

Parameters:

Name Type Description Default
u Float[Array, 'Nz Ny Nx']

x-velocity at U-points.

 required
v Float[Array, 'Nz Ny Nx']

y-velocity at V-points.

 required
cd float or Float[Array, 'Ny Nx']

Dimensionless drag coefficient — a scalar, or a T-point field (ghost ring filled) interpolated to U/V-points.

 required
h_bot float or Float[Array, 'Ny Nx']

Bottom-layer thickness [m] — a scalar, or a T-point field (ghost ring filled) interpolated to U/V-points.

 required

Returns:

Type Description
tuple[Float[Array, 'Nz Ny Nx'], Float[Array, 'Nz Ny Nx']]

(du_drag, dv_drag) at U- and V-points, zero outside k = 1 and in the ghost ring. When self.mask is set, dry faces are zeroed.
Source code in finitevolx/_src/forcing/drag.py 
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
def __call__(
    self,
    u: Float[Array, "Nz Ny Nx"],
    v: Float[Array, "Nz Ny Nx"],
    cd: float | Float[Array, "Ny Nx"],
    h_bot: float | Float[Array, "Ny Nx"],
) -> tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]:
    """Quadratic drag tendencies, non-zero only at the bottom interior level.

    Parameters
    ----------
    u : Float[Array, "Nz Ny Nx"]
        x-velocity at U-points.
    v : Float[Array, "Nz Ny Nx"]
        y-velocity at V-points.
    cd : float or Float[Array, "Ny Nx"]
        Dimensionless drag coefficient — a scalar, or a T-point field
        (ghost ring filled) interpolated to U/V-points.
    h_bot : float or Float[Array, "Ny Nx"]
        Bottom-layer thickness [m] — a scalar, or a T-point field
        (ghost ring filled) interpolated to U/V-points.

    Returns
    -------
    tuple[Float[Array, "Nz Ny Nx"], Float[Array, "Nz Ny Nx"]]
        ``(du_drag, dv_drag)`` at U- and V-points, zero outside
        ``k = 1`` and in the ghost ring.  When ``self.mask`` is
        set, dry faces are zeroed.
    """
    mu = None if self.mask is None else self.mask.u[_K_BOT]
    mv = None if self.mask is None else self.mask.v[_K_BOT]
    du_2d, dv_2d = self._drag2d._tendency(u[_K_BOT], v[_K_BOT], cd, h_bot, mu, mv)
    return _inject_bottom(du_2d, dv_2d, u, v)

finitevolx.RayleighDamping3D

Bases: AbstractForcing

Rayleigh damping of a T-point field on a 3-D Arakawa C-grid.

Applies the :class:RayleighDamping2D stencil uniformly at every z-level and zeroes the z-ghost slices:

dq[k, j, i] = -r[k, j, i] * (q[k, j, i] - q_ref[k, j, i])

r and q_ref may each be a scalar, a per-level profile [Nz], a horizontal field [Ny, Nx] (broadcast over depth), or a full [Nz, Ny, Nx] field.

Parameters:

Name Type Description Default
grid CartesianGrid3D

The underlying 3-D grid.

 required
mask Mask3D or None

Optional land/ocean mask. The inner :class:RayleighDamping2D is mask-free; dry cells of the 3-D result are zeroed with mask.h via jnp.where, so NaN-filled land never leaks into the output.

 None

Examples:

>>> import jax.numpy as jnp
>>> from finitevolx import CartesianGrid3D, RayleighDamping3D
>>> grid = CartesianGrid3D.from_interior(6, 6, 4, 1.0, 1.0, 1.0)
>>> damp = RayleighDamping3D(grid=grid)
>>> q = jnp.ones((grid.Nz, grid.Ny, grid.Nx))
>>> r = jnp.linspace(0.0, 1e-5, grid.Nz)  # per-level damping rate
>>> dq = damp(q, r)
Source code in finitevolx/_src/forcing/rayleigh.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
class RayleighDamping3D(AbstractForcing):
    """Rayleigh damping of a T-point field on a 3-D Arakawa C-grid.

    Applies the :class:`RayleighDamping2D` stencil uniformly at every
    z-level and zeroes the z-ghost slices:

        dq[k, j, i] = -r[k, j, i] * (q[k, j, i] - q_ref[k, j, i])

    ``r`` and ``q_ref`` may each be a scalar, a per-level profile ``[Nz]``,
    a horizontal field ``[Ny, Nx]`` (broadcast over depth), or a full
    ``[Nz, Ny, Nx]`` field.

    Parameters
    ----------
    grid : CartesianGrid3D
        The underlying 3-D grid.
    mask : Mask3D or None, optional
        Optional land/ocean mask.  The inner :class:`RayleighDamping2D` is
        mask-free; dry cells of the 3-D result are zeroed with ``mask.h``
        via ``jnp.where``, so NaN-filled land never leaks into the output.

    Examples
    --------
    >>> import jax.numpy as jnp
    >>> from finitevolx import CartesianGrid3D, RayleighDamping3D
    >>> grid = CartesianGrid3D.from_interior(6, 6, 4, 1.0, 1.0, 1.0)
    >>> damp = RayleighDamping3D(grid=grid)
    >>> q = jnp.ones((grid.Nz, grid.Ny, grid.Nx))
    >>> r = jnp.linspace(0.0, 1e-5, grid.Nz)  # per-level damping rate
    >>> dq = damp(q, r)
    """

    grid: CartesianGrid3D
    mask: Mask3D | None
    _damp2d: RayleighDamping2D

    def __init__(
        self,
        grid: CartesianGrid3D,
        mask: Mask3D | None = None,
    ) -> None:
        self.grid = grid
        self.mask = mask
        # The inner 2-D op is mask-free; the 3-D wrapper owns the mask.
        self._damp2d = RayleighDamping2D(grid=grid.horizontal_grid())

    def __call__(
        self,
        q: Float[Array, "Nz Ny Nx"],
        r: float
        | Float[Array, "Nz"]
        | Float[Array, "Ny Nx"]
        | Float[Array, "Nz Ny Nx"],
        q_ref: float
        | Float[Array, "Nz"]
        | Float[Array, "Ny Nx"]
        | Float[Array, "Nz Ny Nx"]
        | None = None,
    ) -> Float[Array, "Nz Ny Nx"]:
        """Rayleigh damping tendency at T-points over all z-levels.

        Parameters
        ----------
        q : Float[Array, "Nz Ny Nx"]
            Field at T-points.
        r : float or Float[Array, "Nz"] or Float[Array, "Ny Nx"] or Float[Array, "Nz Ny Nx"]
            Damping rate [1/s].
        q_ref : same options as ``r``, or None, optional
            Reference state.  ``None`` (default) damps toward zero.

        Returns
        -------
        Float[Array, "Nz Ny Nx"]
            Damping tendency at T-points, zero in the ghost ring and the
            z-ghost slices.  When ``self.mask`` is set, dry cells are
            zeroed.
        """
        if self.mask is not None:
            # Zero dry cells first: NaN-filled land would otherwise give NaN
            # gradients through -r * (q - q_ref) even where dq is masked.
            q = mask_where(q, self.mask.h)
            if q_ref is not None and jnp.ndim(q_ref) == 3:
                q_ref = mask_where(jnp.asarray(q_ref), self.mask.h)

        r_arr = jnp.asarray(r)
        if q_ref is None:
            dq = eqx.filter_vmap(
                self._damp2d,
                in_axes=(0, _level_axis(r_arr)),
            )(q, r_arr)
        else:
            ref_arr = jnp.asarray(q_ref)
            dq = eqx.filter_vmap(
                self._damp2d,
                in_axes=(0, _level_axis(r_arr), _level_axis(ref_arr)),
            )(q, r_arr, ref_arr)

        # Zero the z-ghost slices to match the 3-D ghost-ring convention.
        dq = zero_z_ghosts(dq)

        if self.mask is None:
            return dq
        return mask_where(dq, self.mask.h)

__call__(q, r, q_ref=None)

Rayleigh damping tendency at T-points over all z-levels.

Parameters:

Name Type Description Default
q Float[Array, 'Nz Ny Nx']

Field at T-points.

 required
r float or Float[Array, 'Nz'] or Float[Array, 'Ny Nx'] or Float[Array, 'Nz Ny Nx']

Damping rate [1/s].

 required
q_ref same options as ``r``, or None

Reference state. None (default) damps toward zero.

 None

Returns:

Type Description
Float[Array, 'Nz Ny Nx']

Damping tendency at T-points, zero in the ghost ring and the z-ghost slices. When self.mask is set, dry cells are zeroed.
Source code in finitevolx/_src/forcing/rayleigh.py 
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
def __call__(
    self,
    q: Float[Array, "Nz Ny Nx"],
    r: float
    | Float[Array, "Nz"]
    | Float[Array, "Ny Nx"]
    | Float[Array, "Nz Ny Nx"],
    q_ref: float
    | Float[Array, "Nz"]
    | Float[Array, "Ny Nx"]
    | Float[Array, "Nz Ny Nx"]
    | None = None,
) -> Float[Array, "Nz Ny Nx"]:
    """Rayleigh damping tendency at T-points over all z-levels.

    Parameters
    ----------
    q : Float[Array, "Nz Ny Nx"]
        Field at T-points.
    r : float or Float[Array, "Nz"] or Float[Array, "Ny Nx"] or Float[Array, "Nz Ny Nx"]
        Damping rate [1/s].
    q_ref : same options as ``r``, or None, optional
        Reference state.  ``None`` (default) damps toward zero.

    Returns
    -------
    Float[Array, "Nz Ny Nx"]
        Damping tendency at T-points, zero in the ghost ring and the
        z-ghost slices.  When ``self.mask`` is set, dry cells are
        zeroed.
    """
    if self.mask is not None:
        # Zero dry cells first: NaN-filled land would otherwise give NaN
        # gradients through -r * (q - q_ref) even where dq is masked.
        q = mask_where(q, self.mask.h)
        if q_ref is not None and jnp.ndim(q_ref) == 3:
            q_ref = mask_where(jnp.asarray(q_ref), self.mask.h)

    r_arr = jnp.asarray(r)
    if q_ref is None:
        dq = eqx.filter_vmap(
            self._damp2d,
            in_axes=(0, _level_axis(r_arr)),
        )(q, r_arr)
    else:
        ref_arr = jnp.asarray(q_ref)
        dq = eqx.filter_vmap(
            self._damp2d,
            in_axes=(0, _level_axis(r_arr), _level_axis(ref_arr)),
        )(q, r_arr, ref_arr)

    # Zero the z-ghost slices to match the 3-D ghost-ring convention.
    dq = zero_z_ghosts(dq)

    if self.mask is None:
        return dq
    return mask_where(dq, self.mask.h)