Module curlew.fields.clebsch

Clebsch (streamfunction) velocity fields built from :class:~curlew.fields.series.FSF potentials.

At fixed lift coordinate w, the velocity is divergence-free by construction:

2-D:  v = ∇⊥β = (∂β/∂x₂, −∂β/∂x₁)
3-D:  v = ∇α × ∇β

Jacobians are analytic (via FSF.gradient_and_hessian) for use by diffeomorphic integrators (forthcoming).

Functions

def clebsch_velocity(dim: int,
*,
n_features: int = 128,
length_scale_range: Tuple[float, float] = (0.5, 3.0),
lift_dim: int = 0,
anchor=None,
killing: str = 'full',
seed: int = 42,
**fsf_kwargs) ‑> ClebschVelocity

Expand source code

def clebsch_velocity(
    dim: int,
    *,
    n_features: int = 128,
    length_scale_range: Tuple[float, float] = (0.5, 3.0),
    lift_dim: int = 0,
    anchor=None,
    killing: str = "full",
    seed: int = 42,
    **fsf_kwargs,
) -> ClebschVelocity:
    """
    Build :class:`FSF` potentials and wrap in :class:`ClebschVelocity`.

    Parameters
    ----------
    dim : int
        Spatial dimension (2 or 3).
    n_features, length_scale_range, lift_dim, seed
        Forwarded to each underlying :class:`FSF`.
    anchor, killing
        Passed to :class:`ClebschVelocity`.
    **fsf_kwargs
        Additional keywords for :class:`FSF` (e.g. ``freq_sampling``, ``scale``).

    Returns
    -------
    ClebschVelocity
    """
    if dim not in (2, 3):
        raise ValueError(f"dim must be 2 or 3, got {dim}")
    xavier = fsf_kwargs.pop("xavier_amplitudes", False)
    h = HSet()
    beta = FSF(
        "beta",
        H=h,
        input_dim=dim,
        lift_dim=lift_dim,
        rff_features=n_features,
        length_scale_range=length_scale_range,
        seed=seed,
        xavier_amplitudes=xavier,
        **fsf_kwargs,
    )
    alpha = None
    if dim == 3:
        alpha = FSF(
            "alpha",
            H=h.copy(),
            input_dim=dim,
            lift_dim=lift_dim,
            rff_features=n_features,
            length_scale_range=length_scale_range,
            seed=seed + 1,
            xavier_amplitudes=xavier,
            **fsf_kwargs,
        )
    return ClebschVelocity(beta, alpha, anchor=anchor, killing=killing)

Build :class:FSF potentials and wrap in :class:ClebschVelocity.

Parameters

dim : int

Spatial dimension (2 or 3).

n_features, length_scale_range, lift_dim, seed

Forwarded to each underlying :class:FSF.

anchor, killing

Passed to :class:ClebschVelocity.

**fsf_kwargs

Additional keywords for :class:FSF (e.g. freq_sampling, scale).

Returns

ClebschVelocity

 

def temporal_clebsch_velocity(dim: int,
*,
n_features: int = 128,
length_scale_range: Tuple[float, float] = (0.5, 3.0),
lift_dim: int = 0,
anchor=None,
killing: str = 'full',
seed: int = 42,
init_std: float = 0.01,
**temporal_fsf_kwargs) ‑> ClebschVelocity

Expand source code

def temporal_clebsch_velocity(
    dim: int,
    *,
    n_features: int = 128,
    length_scale_range: Tuple[float, float] = (0.5, 3.0),
    lift_dim: int = 0,
    anchor=None,
    killing: str = "full",
    seed: int = 42,
    init_std: float = 0.01,
    **temporal_fsf_kwargs,
) -> ClebschVelocity:
    """
    Build :class:`ClebschVelocity` from :class:`~curlew.fields.series.TemporalFSF`
    potentials with a coarse-to-fine spectral schedule over ``t ∈ [0, 1]``.

    Parameters
    ----------
    dim : int
        Spatial dimension (2 or 3).
    n_features, length_scale_range, lift_dim, seed, init_std
        Forwarded to each :class:`TemporalFSF` (``length_scale_range`` is
        ``(λ_fine, λ_coarse)`` for the spectral cutoff sweep).
    anchor, killing
        Passed to :class:`ClebschVelocity`.
    **temporal_fsf_kwargs
        Additional keywords for :class:`TemporalFSF` (e.g. ``freq_sampling``,
        ``activation``, ``scale``, ``freq_damp``, ``spectral_gate``,
        ``spectral_band_width``, ``learnable_directions``).
    """
    if dim not in (2, 3):
        raise ValueError(f"dim must be 2 or 3, got {dim}")
    h = HSet()
    beta = TemporalFSF(
        "beta",
        H=h,
        input_dim=dim,
        lift_dim=lift_dim,
        rff_features=n_features,
        length_scale_range=length_scale_range,
        seed=seed,
        init_std=init_std,
        **temporal_fsf_kwargs,
    )
    alpha = None
    if dim == 3:
        alpha = TemporalFSF(
            "alpha",
            H=h.copy(),
            input_dim=dim,
            lift_dim=lift_dim,
            rff_features=n_features,
            length_scale_range=length_scale_range,
            seed=seed + 1,
            init_std=init_std,
            **temporal_fsf_kwargs,
        )
    return ClebschVelocity(beta, alpha, anchor=anchor, killing=killing)

Build :class:ClebschVelocity from :class:~curlew.fields.series.TemporalFSF potentials with a coarse-to-fine spectral schedule over t ∈ [0, 1].

Parameters

dim : int

Spatial dimension (2 or 3).

n_features, length_scale_range, lift_dim, seed, init_std

Forwarded to each :class:TemporalFSF (length_scale_range is (λ_fine, λ_coarse) for the spectral cutoff sweep).

anchor, killing

Passed to :class:ClebschVelocity.

**temporal_fsf_kwargs

Additional keywords for :class:TemporalFSF (e.g. freq_sampling, activation, scale, freq_damp, spectral_gate, spectral_band_width, learnable_directions).

Classes

class ClebschVelocity (beta: FSF,
alpha: FSF | None = None,
anchor=None,
killing: str = 'full')

Expand source code

class ClebschVelocity(nn.Module):
    """
    Divergence-free velocity from one (2-D) or two (3-D) :class:`FSF` potentials.

    Parameters
    ----------
    beta : FSF
        Streamfunction potential β. In 2-D, v = ∇⊥β; in 3-D, v = ∇α × ∇β.
    alpha : FSF, optional
        Second potential for 3-D; required when ``beta.dim == 3``.
    anchor : array-like, optional
        Point where rigid modes are removed (see ``killing``).
    killing : ``"full"`` | ``"translation"``
        ``"full"`` removes translation and rotation; ``"translation"`` only
        removes uniform translation.

    Notes
    -----
    Potentials are evaluated on the **raw** FSF path (``gradient_and_hessian``),
    without geomodelling ``Transform`` or drift. For lifted potentials (fault
    sheet labels), pass ``w`` to ``forward``; derivatives are w.r.t. spatial
    coordinates only. Pass integration time ``t`` when the underlying
    potentials are :class:`~curlew.fields.series.TemporalFSF` instances.
    """

    def __init__(
        self,
        beta: FSF,
        alpha: Optional[FSF] = None,
        anchor=None,
        killing: str = "full",
    ) -> None:
        super().__init__()
        if beta.dim not in (2, 3):
            raise ValueError(f"beta.dim must be 2 or 3, got {beta.dim}")
        if beta.dim == 3 and alpha is None:
            raise ValueError("alpha potential is required for dim=3")
        if alpha is not None and alpha.dim != beta.dim:
            raise ValueError("alpha and beta must have the same spatial dimension")
        if killing not in ("full", "translation"):
            raise ValueError(f"unknown killing mode {killing!r}")

        self.dim = beta.dim
        self.beta = beta
        self.alpha = alpha
        self.killing = killing

        if anchor is not None:
            a = _tensor(anchor, dev=curlew.device, dt=curlew.dtype).reshape(1, self.dim)
        else:
            if hasattr(beta, "A_mu"):
                amp = beta.A_mu
            elif hasattr(beta, "A_rho"):
                amp = beta.A_rho
            else:
                amp = beta.A
            a = torch.zeros(0, self.dim, dtype=amp.dtype, device=amp.device)
        self.register_buffer("anchor", a)

    @property
    def anchored(self) -> bool:
        return self.anchor.numel() > 0

    def _raw(
        self,
        x: torch.Tensor,
        w: Optional[torch.Tensor],
        jac: bool,
        t: Optional[float] = None,
    ):
        g_b, h_b = _pot_gradient_and_hessian(self.beta, x, w, t)
        if self.dim == 2:
            v = torch.stack([g_b[:, 1], -g_b[:, 0]], dim=1)
            if not jac:
                return v, None
            jv = torch.stack([h_b[:, :, 1], -h_b[:, :, 0]], dim=1)
            return v, jv

        g_a, h_a = _pot_gradient_and_hessian(self.alpha, x, w, t)
        v = torch.linalg.cross(g_a, g_b, dim=-1)
        if not jac:
            return v, None
        jv = (
            torch.linalg.cross(h_a, g_b.unsqueeze(1).expand_as(h_a), dim=-1).mT
            + torch.linalg.cross(g_a.unsqueeze(1).expand_as(h_b), h_b, dim=-1).mT
        )
        return v, jv

    def forward(
        self,
        x: torch.Tensor,
        w: Optional[torch.Tensor] = None,
        t: Optional[float] = None,
    ) -> torch.Tensor:
        """Velocity v(x[, w], t) with shape (N, dim)."""
        v, _ = self._raw(x, w, jac=False, t=t)
        if self.anchored:
            k, _ = self._gauge(x, t=t)
            v = v - k
        return v

    def forward_and_jacobian(
        self,
        x: torch.Tensor,
        w: Optional[torch.Tensor] = None,
        t: Optional[float] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Return velocity (N, dim) and Jacobian ∂v/∂x (N, dim, dim)."""
        v, jv = self._raw(x, w, jac=True, t=t)
        if self.anchored:
            k, jk = self._gauge(x, t=t)
            v, jv = v - k, jv - jk
        return v, jv

    def _gauge(self, x: torch.Tensor, t: Optional[float] = None):
        v0, j0 = self._raw(self.anchor, None, jac=self.killing == "full", t=t)
        j0 = j0[0] if j0 is not None else None
        return _killing_frame(v0, j0, x, self.anchor, self.killing)

Divergence-free velocity from one (2-D) or two (3-D) :class:FSF potentials.

Parameters

beta : FSF

Streamfunction potential β. In 2-D, v = ∇⊥β; in 3-D, v = ∇α × ∇β.

alpha : FSF, optional

Second potential for 3-D; required when beta.dim == 3.

anchor : array-like, optional

Point where rigid modes are removed (see killing).

killing : "full"`` | ``"translation"

"full" removes translation and rotation; "translation" only removes uniform translation.

Notes

Potentials are evaluated on the raw FSF path (gradient_and_hessian), without geomodelling Transform or drift. For lifted potentials (fault sheet labels), pass w to forward; derivatives are w.r.t. spatial coordinates only. Pass integration time t when the underlying potentials are :class:~curlew.fields.series.TemporalFSF instances.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Ancestors

• torch.nn.modules.module.Module

Instance variables

prop anchored : bool

Expand source code

@property
def anchored(self) -> bool:
    return self.anchor.numel() > 0

Methods

def forward(self, x: torch.Tensor, w: torch.Tensor | None = None, t: float | None = None) ‑> torch.Tensor

Expand source code

def forward(
    self,
    x: torch.Tensor,
    w: Optional[torch.Tensor] = None,
    t: Optional[float] = None,
) -> torch.Tensor:
    """Velocity v(x[, w], t) with shape (N, dim)."""
    v, _ = self._raw(x, w, jac=False, t=t)
    if self.anchored:
        k, _ = self._gauge(x, t=t)
        v = v - k
    return v

Velocity v(x[, w], t) with shape (N, dim).

def forward_and_jacobian(self, x: torch.Tensor, w: torch.Tensor | None = None, t: float | None = None) ‑> Tuple[torch.Tensor, torch.Tensor]

Expand source code

def forward_and_jacobian(
    self,
    x: torch.Tensor,
    w: Optional[torch.Tensor] = None,
    t: Optional[float] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Return velocity (N, dim) and Jacobian ∂v/∂x (N, dim, dim)."""
    v, jv = self._raw(x, w, jac=True, t=t)
    if self.anchored:
        k, jk = self._gauge(x, t=t)
        v, jv = v - k, jv - jk
    return v, jv

Return velocity (N, dim) and Jacobian ∂v/∂x (N, dim, dim).