funcflows 0.1.2__py3-none-any.whl

FuncFlows/transports/abstract_transformation.py

@@ -0,0 +1,32 @@
+import torch
+
+from FuncFlows.base_measures import ReferenceMeasure
+
+
+class Transformation(torch.nn.Module):
+    """A measurable bijection of coefficient space, with its Radon-Nikodym derivative.
+
+    `context` threads through every method so conditional fields work end to end; unconditional
+    transformations simply ignore it.
+    """
+
+    def __init__(self, base_measure: ReferenceMeasure):
+        super().__init__()
+        self.base_measure = base_measure
+
+    def _map(self, coeffs, context=None):
+        """coeffs [..., num_functions] -> (coeffs_out [..., num_functions], log_det_term [...])
+        log_det_term = log|det Df| for a discrete map, integral of Tr Dh for an ODE."""
+        raise NotImplementedError
+
+    def pull_back(self, coeffs_out, context=None):
+        """Inverse map. Optional; needed only by NLL / learned-prior uses."""
+        raise NotImplementedError
+
+    def push_forward(self, coeffs, context=None):
+        """-> (coeffs_out, log_rn_weight) with log_rn_weight = log d(f#mu)/d(mu) at coeffs_out."""
+        coeffs_out, log_det_term = self._map(coeffs, context)
+        return coeffs_out, self.base_measure.log_density_diff(coeffs, coeffs_out) - log_det_term
+
+    def log_rn_at(self, coeffs_out, context=None):
+        return self.push_forward(self.pull_back(coeffs_out, context), context)[1]

FuncFlows/transports/continuous/__init__.py

@@ -0,0 +1,6 @@
+from .vector_fields import (LinearField, MatrixField,
+                            SumField, VectorField, TimeCosines)
+from .base_continuous import ContinuousTransformation
+from .conditioners import (Conditioner, TimeBasisConditioner,
+                           DataConditioner)
+from.grid_fields import (PointwiseField, OperatorField, GridTransform)

FuncFlows/transports/continuous/base_continuous.py

@@ -0,0 +1,63 @@
+import torch
+
+from ..abstract_transformation import Transformation
+
+
+class ContinuousTransformation(Transformation):
+    """f = the time-T flow of dv/dt = h(v,t). log_det_term is the integral of Tr Dh (P2 Thm 5).
+
+    One integrator serves every use. with_trace=False skips the log-determinant entirely, which
+    is what transport (sampling) and pull_back need; with_trace=True uses velocity_and_trace so
+    the trace shares the field evaluation instead of repeating it. The old version recomputed the
+    field twice per RK4 stage in every density-bearing solve.
+    """
+
+    def __init__(self, base_measure, vector_field, total_time=1.0, num_steps=20, method="rk4"):
+        super().__init__(base_measure)
+        self.vector_field, self.total_time, self.num_steps, self.method = vector_field, total_time, num_steps, method
+
+    def _stage(self, coeffs, time_value, context, with_trace):
+        if with_trace:
+            return self.vector_field.velocity_and_trace(coeffs, time_value, context)
+        return self.vector_field.velocity(coeffs, time_value, context), None
+
+    def _integrate(self, coeffs, forward=True, context=None, with_trace=True):
+        step = (self.total_time / self.num_steps) * (1 if forward else -1)
+        time_value = 0.0 if forward else self.total_time
+        trace_total = torch.zeros_like(coeffs[..., 0]) if with_trace else None
+        for _ in range(self.num_steps):
+            velocity_1, trace_1 = self._stage(coeffs, time_value, context, with_trace)
+            if self.method == "euler":
+                velocity, trace = velocity_1, trace_1
+            else:
+                velocity_2, trace_2 = self._stage(coeffs + 0.5 * step * velocity_1,
+                                                  time_value + 0.5 * step, context, with_trace)
+                velocity_3, trace_3 = self._stage(coeffs + 0.5 * step * velocity_2,
+                                                  time_value + 0.5 * step, context, with_trace)
+                velocity_4, trace_4 = self._stage(coeffs + step * velocity_3,
+                                                  time_value + step, context, with_trace)
+                velocity = (velocity_1 + 2 * velocity_2 + 2 * velocity_3 + velocity_4) / 6
+                trace = ((trace_1 + 2 * trace_2 + 2 * trace_3 + trace_4) / 6) if with_trace else None
+            coeffs = coeffs + step * velocity
+            if with_trace:
+                trace_total = trace_total + step * trace
+            time_value = time_value + step
+        return coeffs, trace_total
+
+    def _map(self, coeffs, context=None):
+        return self._integrate(coeffs, forward=True, context=context)
+
+    def pull_back(self, coeffs_out, context=None):
+        return self._integrate(coeffs_out, forward=False, context=context, with_trace=False)[0]
+
+    def log_rn_at(self, coeffs_out, context=None):
+        """One backward solve instead of pull_back + push_forward (P2 Supp Thm 2.1)."""
+        coeffs, trace_total = self._integrate(coeffs_out, forward=False, context=context)
+        return self.base_measure.log_density_diff(coeffs, coeffs_out) + trace_total
+
+    def transport(self, coeffs, context=None):
+        """Integrate the field only. No trace, no log-determinant — for sampling."""
+        return self._integrate(coeffs, forward=True, context=context, with_trace=False)[0]
+
+
+

FuncFlows/transports/continuous/conditioners.py

@@ -0,0 +1,83 @@
+import torch
+
+from .vector_fields import TimeCosines
+
+# base class
+class Conditioner(torch.nn.Module):
+    """(time, context) -> (output_matrix [num_functions, num_terms],
+                           input_matrix  [num_terms, num_functions],
+                           biases        [num_terms] or [batch, num_terms])"""
+
+
+    def __init__(self, num_functions, num_terms):
+        super().__init__()
+        self.num_functions, self.num_terms = num_functions, num_terms
+
+
+    def forward(self, time_value, context=None):
+        raise NotImplementedError
+
+
+
+# Just a class for conditioners in the velocity fields script
+    # that handles unknown parameters that are functions of just time
+class TimeBasisConditioner(Conditioner):
+    """A(t) = sum_m cos(m pi t / T) A_m, same for B and the biases. Genuine time dependence
+    for num_time_modes * 2 * L * r parameters, instead of a net emitting that many per step.
+
+    The contraction over m is a tensordot, which never materialises the [T, r, L] product the
+    broadcast-and-sum form did — that intermediate was rebuilt four times per RK4 step.
+    """
+
+    def __init__(self, num_functions, num_terms, num_time_modes=4, total_time=1.0,
+                 init_scale=0.05, dtype=torch.float64):
+        super().__init__(num_functions, num_terms)
+        scale = init_scale / (num_functions * num_time_modes) ** 0.5
+        self.total_time, self.num_time_modes = total_time, num_time_modes
+
+        self.output_stack = torch.nn.Parameter(scale * torch.randn(num_time_modes, num_functions, num_terms, dtype=dtype))
+        self.input_stack = torch.nn.Parameter(scale * torch.randn(num_time_modes, num_terms, num_functions, dtype=dtype))
+        self.bias_stack = torch.nn.Parameter(torch.zeros(num_time_modes, num_terms, dtype=dtype))
+
+        self.cosines = TimeCosines(num_time_modes, total_time, dtype)     # cached per scalar time
+
+
+    def time_weights(self, time_value):
+        return self.cosines(time_value)
+
+    def forward(self, time_value, context=None):
+        weights = self.time_weights(time_value)
+        return (torch.tensordot(weights, self.output_stack, dims=1),
+                torch.tensordot(weights, self.input_stack, dims=1),
+                weights @ self.bias_stack)
+
+
+
+# Just a class for conditioners in the velocity fields script
+    # that handles unknown parameters that are functions of time AND data embeddings
+class DataConditioner(TimeBasisConditioner):
+    """Same field, but the data shifts the bias: b(t, c) = sum_m w_m(t) [bias_m + c @ context_m.T].
+
+    The context does not depend on v, so dh/dv is unchanged and MatrixField.trace still holds.
+    """
+
+    def __init__(self, num_functions, num_terms, context_dim, num_time_modes=4,
+                 total_time=1.0, init_scale=0.05, dtype=torch.float64):
+        super().__init__(num_functions, num_terms, num_time_modes, total_time, init_scale, dtype)
+
+        self.context_stack = torch.nn.Parameter(
+            init_scale / (context_dim * num_time_modes) ** 0.5
+            * torch.randn(num_time_modes, num_terms, context_dim, dtype=dtype))
+
+
+
+    def forward(self, time_value, context=None):
+        output_matrix, input_matrix, biases = super().forward(time_value)
+
+
+        if context is not None:
+            weights = self.time_weights(time_value)
+            biases = biases + context @ torch.tensordot(weights, self.context_stack, dims=1).T
+
+
+        return output_matrix, input_matrix, biases

FuncFlows/transports/continuous/grid_fields.py

@@ -0,0 +1,403 @@
+
+import torch
+from .vector_fields import VectorField, TimeCosines, combine_context
+
+
+################################################################################################################
+################################################################################################################
+################################################################################################################
+
+
+class GridTransform(torch.nn.Module):
+    """Coefficients <-> values on a uniform grid, and back, for one fixed basis truncation."""
+
+    def __init__(self, basis, num_active, grid_size, mode="auto", dtype=torch.float64):
+        super().__init__()
+
+        self.num_active, self.grid_size, self.dtype = num_active, grid_size, dtype
+        self.physical_dim = basis.physical_dim
+        self.quadrature_weight = grid_size ** (-basis.physical_dim)
+        device = basis.laplacian_eigenvalues.device
+
+
+        # cell-edge grid, so the FFT needs no midpoint phase; on a periodic uniform grid the
+        # rectangle rule is spectrally accurate either way
+        axis = torch.arange(grid_size, dtype=dtype, device=device) / grid_size
+        points = torch.cartesian_prod(*[axis] * basis.physical_dim).reshape(-1, basis.physical_dim)
+        values = basis.evaluate(points)[:, :num_active]
+
+        self.register_buffer("basis_grid", values)                 # [grid^d, num_active]
+
+        self.use_fft = False
+
+        if mode != "dense" and hasattr(basis, "wavenumbers"):      # separable Fourier basis only
+            self._build_tables(basis, num_active, grid_size, device)
+            self.use_fft = self._agrees(values)
+
+
+        if mode == "fft" and not self.use_fft:
+            raise RuntimeError("the FFT transform needs a FourierBasis and must pass its check; "
+                               "pass transform='dense'")
+        
+        self.register_buffer("basis_values", None if self.use_fft else values)
+
+
+
+    def _build_tables(self, basis, num_active, grid_size, device):
+        """Per-axis mode index m: 0 -> constant, odd -> cos with k=(m+1)//2, even>0 -> sin with k=m//2."""
+        wavenumbers = basis.wavenumbers[:num_active].round().long()        # [num_active, d]
+
+        is_sine = basis.phases[:num_active] != 0
+
+        modes = torch.where(wavenumbers == 0, torch.zeros_like(wavenumbers),
+                            2 * wavenumbers - 1 + is_sine.long())
+        
+        self.max_axis_mode = int(modes.max().item()) + 1
+        highest = int(wavenumbers.max().item())
+
+
+        if highest >= grid_size // 2:
+            raise ValueError(f"grid_size {grid_size} is below Nyquist for wavenumber {highest}; "
+                             f"use at least {2 * highest + 2}")
+        flat = torch.zeros(num_active, dtype=torch.long, device=device)
+
+
+        for axis in range(self.physical_dim):
+            flat = flat * self.max_axis_mode + modes[:, axis]
+
+
+        
+        self.register_buffer("flat_index", flat)
+        axis_modes = torch.arange(1, self.max_axis_mode, device=device)
+        cosine, sine = axis_modes[axis_modes % 2 == 1], axis_modes[axis_modes % 2 == 0]
+
+
+        self.register_buffer("cosine_mode", cosine)
+        self.register_buffer("cosine_wave", (cosine + 1) // 2)
+        self.register_buffer("sine_mode", sine)
+        self.register_buffer("sine_wave", sine // 2)
+        self.register_buffer("zero_index", torch.zeros(1, dtype=torch.long, device=device))
+
+
+
+    def _synthesise_axis(self, coefficients):
+        """[..., max_axis_mode] -> [..., grid_size] along the last axis."""
+        shape = coefficients.shape[:-1] + (self.grid_size // 2 + 1,)
+        scale = self.grid_size / 2 ** 0.5
+
+        real = torch.zeros(shape, dtype=coefficients.dtype, device=coefficients.device)
+        real = real.index_add(-1, self.zero_index, self.grid_size * coefficients[..., :1])
+        real = real.index_add(-1, self.cosine_wave, scale * coefficients[..., self.cosine_mode])
+
+        imaginary = torch.zeros(shape, dtype=coefficients.dtype, device=coefficients.device)
+        imaginary = imaginary.index_add(-1, self.sine_wave, -scale * coefficients[..., self.sine_mode])
+
+        return torch.fft.irfft(torch.complex(real, imaginary), n=self.grid_size, dim=-1)
+
+    def _analyse_axis(self, values):
+        spectrum = torch.fft.rfft(values, dim=-1)
+        scale = 2 ** 0.5 / self.grid_size
+
+        coefficients = torch.zeros(values.shape[:-1] + (self.max_axis_mode,),
+                                   dtype=values.dtype, device=values.device)
+        
+        coefficients = coefficients.index_add(-1, self.zero_index, spectrum.real[..., :1] / self.grid_size)
+        coefficients = coefficients.index_add(-1, self.cosine_mode, scale * spectrum.real[..., self.cosine_wave])
+
+        return coefficients.index_add(-1, self.sine_mode, -scale * spectrum.imag[..., self.sine_wave])
+
+    def to_grid(self, coeffs):
+        """[..., num_active] -> [..., grid_size^d]."""
+        if not self.use_fft:
+            return coeffs @ self.basis_values.T
+        
+        block = torch.zeros(coeffs.shape[:-1] + (self.max_axis_mode ** self.physical_dim,),
+                            dtype=coeffs.dtype, device=coeffs.device)
+        block = block.index_add(-1, self.flat_index, coeffs)
+        block = block.reshape(coeffs.shape[:-1] + (self.max_axis_mode,) * self.physical_dim)
+
+        for _ in range(self.physical_dim):
+            block = self._synthesise_axis(block).movedim(-1, -self.physical_dim)
+
+        return block.reshape(coeffs.shape[:-1] + (self.grid_size ** self.physical_dim,))
+
+    def from_grid(self, values):
+        """[..., grid_size^d] -> [..., num_active] projection."""
+        if not self.use_fft:
+            return (values @ self.basis_values) * self.quadrature_weight
+        
+        block = values.reshape(values.shape[:-1] + (self.grid_size,) * self.physical_dim)
+
+        for _ in range(self.physical_dim):
+            block = self._analyse_axis(block).movedim(-1, -self.physical_dim)
+
+        return block.reshape(values.shape[:-1]
+                             + (self.max_axis_mode ** self.physical_dim,))[..., self.flat_index]
+
+    def _agrees(self, values, tolerance=None):
+        """The FFT path is a fixed reindexing of the dense one, so it should agree to round-off."""
+        tolerance = tolerance or (1e-4 if self.dtype is torch.float32 else 1e-9)
+
+        probe = torch.randn(3, self.num_active, dtype=self.dtype, device=values.device)
+        grid_probe = torch.randn(3, values.shape[0], dtype=self.dtype, device=values.device)
+        self.use_fft = True
+
+        try:
+            forward = (self.to_grid(probe) - probe @ values.T).abs().max().item()
+            backward = (self.from_grid(grid_probe)
+                        - (grid_probe @ values) * self.quadrature_weight).abs().max().item()
+        except Exception:
+            self.use_fft = False
+            return False
+        
+        self.use_fft = False
+        scale = max(1.0, probe.abs().max().item())
+        return bool(forward < tolerance * scale and backward < tolerance * scale)
+
+
+
+
+
+
+################################################################################################################
+################################################################################################################
+################################################################################################################
+
+
+class PointwiseField(VectorField):
+    """h(v,t)(x) = a(x) tanh( g(x) v(x) + (K_t v)(x) + c(x) ) — one Fourier-neural-operator layer.
+
+    Taken from arXiv:2010.08895
+    
+    """
+
+    supports_batched_time = True
+
+    def __init__(self, basis, num_active, grid_size, num_time_modes=4, num_field_modes=32,
+                 num_spectral_modes=0, context_dim=0, total_time=1.0, mode_scale=None,
+                 init_scale=0.05, transform="auto", time_pair=False, dtype=torch.float64):
+        super().__init__()
+        self.num_active, self.num_time_modes = num_active, num_time_modes
+        self.total_time, self.mode_scale, self.dtype = total_time, mode_scale, dtype
+        self.physical_dim, self.grid_size = basis.physical_dim, grid_size
+        self.quadrature_weight = grid_size ** (-basis.physical_dim)
+
+        self.cosines = TimeCosines(num_time_modes, total_time, dtype, time_pair)
+
+        self.transform = GridTransform(basis, num_active, grid_size, mode=transform, dtype=dtype)
+
+        self.quadrature_weight = self.transform.quadrature_weight
+        device = basis.laplacian_eigenvalues.device
+        values = self.transform.basis_grid
+
+        self.register_buffer("squared_sum", (values ** 2).sum(1))          # S(x), [grid^d]
+        self.register_buffer("field_values", values[:, :num_field_modes])  # a, g, c stay dense
+
+        gain_scale = init_scale / num_field_modes ** 0.5
+        self.gain_stack = torch.nn.Parameter(
+            gain_scale * torch.randn(num_time_modes, num_field_modes, dtype=dtype))
+        
+        self.slope_stack = torch.nn.Parameter(torch.zeros(num_time_modes, num_field_modes, dtype=dtype))
+        self.shift_stack = torch.nn.Parameter(torch.zeros(num_time_modes, num_field_modes, dtype=dtype))
+
+        with torch.no_grad():                       # phi_0 = 1, so this makes g(x) a constant.
+            # whitened coefficients are O(1) each, so the whitened field is O(sqrt(num_active))
+            self.slope_stack[0, 0] = 1.0 if mode_scale is None else num_active ** -0.5
+
+        self.context_stack = None
+        if context_dim:
+            self.context_stack = torch.nn.Parameter(
+                init_scale / context_dim ** 0.5
+                * torch.randn(num_time_modes, num_field_modes, context_dim, dtype=dtype))
+
+
+        self.kappa_stack = None
+        if num_spectral_modes:
+            logs = torch.log1p(basis.laplacian_eigenvalues[:num_active].to(dtype))
+            span = (logs.max() - logs.min()).clamp(min=1e-12)
+            position = ((logs - logs.min()) / span).clamp(0, 1)
+            orders = torch.arange(num_spectral_modes, dtype=dtype, device=device)
+            self.register_buffer("spectral_features", torch.cos(torch.pi * orders * position[:, None]))
+            self.register_buffer("squared_basis", values ** 2)             # [grid^d, num_active]
+            # zero, not init_scale * randn: the slope path is normalised by num_active^-1/2 so the
+            # pre-activation is O(1), but kappa multiplies whitened coefficients that sum over all
+            # num_active modes. At 0.05 * randn that term is O(0.17 * sqrt(2000)) ~ 7.6, tanh is
+            # flat, (1 - hidden^2) ~ 0, and neither the trace nor any gradient gets through.
+            self.kappa_stack = torch.nn.Parameter(
+                torch.zeros(num_time_modes, num_spectral_modes, dtype=dtype))
+
+
+    ################################################################################
+    ################################################################################
+    # field
+
+    def _profiles(self, time_value, context):
+        weights = self.cosines(time_value)
+        gain = (weights @ self.gain_stack) @ self.field_values.T
+        slope = (weights @ self.slope_stack) @ self.field_values.T
+        shift_coeffs = weights @ self.shift_stack
+
+
+        if context is not None and self.context_stack is not None:
+            shift_coeffs = shift_coeffs + combine_context(self.context_stack, context, weights)
+        
+        multiplier = None
+        if self.kappa_stack is not None:
+            multiplier = (weights @ self.kappa_stack) @ self.spectral_features.T
+        
+        return gain, slope, shift_coeffs @ self.field_values.T, multiplier
+
+
+
+    def _hidden(self, coeffs, time_value, context):
+        gain, slope, shift, multiplier = self._profiles(time_value, context)
+        active = coeffs[..., :self.num_active]
+
+        if self.mode_scale is not None:
+            active = active / self.mode_scale[:self.num_active]
+        
+        if multiplier is None:
+            pre_activation = slope * self.transform.to_grid(active) + shift
+        else:                                       # one transform for both fields, not two
+            both = self.transform.to_grid(torch.stack([active, active * multiplier]))
+            pre_activation = slope * both[0] + both[1] + shift                      # + (K v)(x)
+        
+        return gain, slope, multiplier, torch.tanh(pre_activation)
+
+
+
+    def _update(self, coeffs, gain, hidden):
+        update = self.transform.from_grid(gain * hidden)
+
+        if self.mode_scale is not None:
+            update = update * self.mode_scale[:self.num_active]
+        
+        return torch.nn.functional.pad(update, (0, coeffs.shape[-1] - self.num_active))
+
+
+
+    def _trace(self, gain, slope, multiplier, hidden):
+        local = slope * self.squared_sum
+        if multiplier is not None:
+            local = local + multiplier @ self.squared_basis.T         # T(x) = sum_j kappa_j phi_j^2
+        return ((gain * local) * (1 - hidden ** 2)).sum(-1) * self.quadrature_weight
+
+
+
+    @property
+    def use_fft(self):
+        """Lives on the transform now. Kept here so callers written before it was factored out
+        (the scripts' preflight, among others) don't explode."""
+        return self.transform.use_fft
+
+
+
+    def velocity(self, coeffs, time_value, context=None):
+
+        gain, _, _, hidden = self._hidden(coeffs, time_value, context)
+
+        return self._update(coeffs, gain, hidden)
+
+
+
+    def trace(self, coeffs, time_value, context=None):
+        gain, slope, multiplier, hidden = self._hidden(coeffs, time_value, context)
+
+        return self._trace(gain, slope, multiplier, hidden)
+
+
+
+    def velocity_and_trace(self, coeffs, time_value, context=None):
+
+        gain, slope, multiplier, hidden = self._hidden(coeffs, time_value, context)
+
+        return self._update(coeffs, gain, hidden), self._trace(gain, slope, multiplier, hidden)
+
+
+
+
+
+
+
+
+
+################################################################################################################
+################################################################################################################
+################################################################################################################
+
+
+class OperatorField(VectorField):
+    """A stack of Fourier-neural-operator layers as the velocity (FOT-CFM backbone). 
+
+    Took from arXiv:1810.01367
+
+        z_0     = lift * v(x)
+        z_(l+1) = tanh( W_l z_l + K_l z_l + b_l )      W_l pointwise, K_l mode-diagonal
+        h       = project . z_L,  projected back onto the basis
+    """
+
+    supports_batched_time = True
+
+    def __init__(self, basis, num_active, grid_size, num_channels=4, num_layers=3,
+                 num_time_modes=4, num_spectral_modes=8, context_dim=0, total_time=1.0,
+                 mode_scale=None, init_scale=0.5, transform="auto", time_pair=False,
+                 trace_samples=1, dtype=torch.float64):
+        super().__init__()
+        self.num_active, self.num_layers, self.num_channels = num_active, num_layers, num_channels
+        self.total_time, self.mode_scale, self.dtype = total_time, mode_scale, dtype
+        self.trace_samples = trace_samples
+        self.cosines = TimeCosines(num_time_modes, total_time, dtype, time_pair)
+        self.transform = GridTransform(basis, num_active, grid_size, mode=transform, dtype=dtype)
+
+        logs = torch.log1p(basis.laplacian_eigenvalues[:num_active].to(dtype))
+        span = (logs.max() - logs.min()).clamp(min=1e-12)
+        position = ((logs - logs.min()) / span).clamp(0, 1)
+        orders = torch.arange(num_spectral_modes, dtype=dtype, device=logs.device)
+        # smooth in log wavenumber, so the multipliers do not depend on where the basis is cut
+        self.register_buffer("spectral_features", torch.cos(torch.pi * orders * position[:, None]))
+        # a whitened v(x) is O(sqrt(num_active)), so the lift has to undo that or every tanh
+        # starts saturated and nothing propagates (the same trap as the kappa init)
+        self.register_buffer("lift", torch.full((num_channels,),
+                                                1.0 if mode_scale is None else num_active ** -0.5,
+                                                dtype=dtype))
+        shape = (num_layers, num_time_modes, num_channels, num_channels)
+        self.pointwise_stack = torch.nn.Parameter(
+            init_scale / (num_channels * num_time_modes) ** 0.5 * torch.randn(shape, dtype=dtype))
+        self.spectral_stack = torch.nn.Parameter(torch.zeros(shape + (num_spectral_modes,), dtype=dtype))
+        self.bias_stack = torch.nn.Parameter(
+            torch.zeros(num_layers, num_time_modes, num_channels, dtype=dtype))
+        self.project_stack = torch.nn.Parameter(
+            init_scale / (num_channels * num_time_modes) ** 0.5
+            * torch.randn(num_time_modes, num_channels, dtype=dtype))
+        self.context_stack = None
+        if context_dim:
+            self.context_stack = torch.nn.Parameter(
+                init_scale / (context_dim * num_time_modes) ** 0.5
+                * torch.randn(num_layers, num_time_modes, num_channels, context_dim, dtype=dtype))
+
+    def velocity(self, coeffs, time_value, context=None):
+        weights = self.cosines(time_value)
+        active = coeffs[..., :self.num_active]
+        if self.mode_scale is not None:
+            active = active / self.mode_scale[:self.num_active]
+        state = self.transform.to_grid(active).unsqueeze(-1) * self.lift       # [..., points, channels]
+        for layer in range(self.num_layers):
+            pointwise = torch.tensordot(weights, self.pointwise_stack[layer], dims=1)
+            multipliers = (torch.tensordot(weights, self.spectral_stack[layer], dims=1)
+                           @ self.spectral_features.T)              # [..., out, in, num_active]
+            modes = self.transform.from_grid(state.transpose(-1, -2))
+            spectral = self.transform.to_grid((multipliers * modes.unsqueeze(-3)).sum(-2))
+            biases = torch.tensordot(weights, self.bias_stack[layer], dims=1)
+            if context is not None and self.context_stack is not None:
+                biases = biases + combine_context(self.context_stack[layer], context, weights)
+            state = torch.tanh(torch.matmul(state, pointwise.transpose(-1, -2))
+                               + spectral.transpose(-1, -2) + biases.unsqueeze(-2))
+        projection = (weights @ self.project_stack).unsqueeze(-2)
+        update = self.transform.from_grid((state * projection).sum(-1))
+        if self.mode_scale is not None:
+            update = update * self.mode_scale[:self.num_active]
+        return torch.nn.functional.pad(update, (0, coeffs.shape[-1] - self.num_active))
+
+    def velocity_and_trace(self, coeffs, time_value, context=None):
+        return self.estimated_velocity_and_trace(coeffs, time_value, context)
+