SFI.integrate.timeops module

class SFI.integrate.timeops.TimeOp(fn, name=None, *, batch_safe=False, requires=None)[source]

Bases: object

A time-local operator evaluated on one time slice.

A TimeOp declares which streamed fields it needs either via its function signature or the explicit requires parameter.

Example

>>> @timeop
... def vel(dX_minus, dt_minus):
...     return dX_minus / dt_minus[..., None]
>>> vel.requires
frozenset({'dX_minus', 'dt_minus'})
Parameters:

  • fn (Callable[..., jnp.ndarray])

  • name (str | None)

  • batch_safe (bool)

  • requires (Optional[FrozenSet[str]])

property name: str
property requires: FrozenSet[str]
SFI.integrate.timeops.add(ops, *, name=None)[source]

Return a TimeOp that sums the outputs of all ops element-wise.

Parameters:

  • ops (iterable of TimeOp) – At least one operand is required.

  • name (str | None)

Return type:

TimeOp

SFI.integrate.timeops.scale(op, alpha, *, name=None)[source]

Return a TimeOp that multiplies op by scalar alpha.

Parameters:

  • op (TimeOp)

  • alpha (float | int)

  • name (str | None)

Return type:

TimeOp

SFI.integrate.timeops.stream(key, *, name=None)[source]

Return a TimeOp that passes stream key through unchanged.

Parameters:

  • key (str)

  • name (str | None)

Return type:

TimeOp

SFI.integrate.timeops.timeop(fn=None, *, name=None, batch_safe=False, requires=None)[source]

Decorator: convert a function into a TimeOp.

Parameters:

  • batch_safe (bool) – If True, the function already handles a leading batch (K) axis in its inputs without requiring an additional jax.vmap.

  • requires (frozenset, optional) – Explicit set of required stream keys. When given, the function signature is not inspected for parameter names. Use this when the function accepts **streams and the required keys cannot be inferred from the signature.

  • fn (Callable[[...], Array] | None)

  • name (str | None)

Notes

Without requires, the function’s parameter names define the required stream keys (**kwargs/*args parameters are ignored).

SFI.integrate.timeops.velocity(dx_key, dt_key, *, name=None)[source]

Return a TimeOp that computes dx/dt, broadcasting dt over arbitrary leading dims.

Parameters:

  • dx_key (str)

  • dt_key (str)

  • name (str | None)

Return type:

TimeOp