jaxdf.operators.functions

Operators

compose

Implements a decorator that allows to compose jax functions with fields. Given a function $f$ and a Field $x$, the result is a new field representing

$$ y = f(x) $$

The usage of the decorator is as follows:

y = compose(x)(f)

Concrete implementations:

compose(x: 'Continuous', *, params = None)

Applies function composition on the get_fun of the Continuous object.

compose(x: 'OnGrid', *, params = None)

Maps the given function over the pytree of parameters of the Field.

compose(x: 'object', *, params = None)

For non-field objects, the composition is simply the application of the jax function to the input.

compose(x)(fun) == fun(x)

functional

It works exactly like compose, but the function f maps one or more fields to a scalar value.

The usage of the decorator is as follows:

y = functional(x)(f)

It is useful to improve the readibility of the code.

For non-field objects, a functional is simply the application of the jax function to the input.

functional(x)(fun) == fun(x)

Concrete implementations:

functional(x: 'OnGrid', *, params = None)

Maps the given function over the parameters of the field

x = OnGrid(params=-1.0, ...)
y = functional(x)(jnp.sum)
y.params # This is -1.0

functional(x: 'object', *, params = None)

For non-field objects, a functional is simply the application of the jax function to the input.

functional(x)(fun) == fun(x)

get_component

This operator $A(u, \text{dim})$ which has the signature (u: Field, dim: int) -> Field. It returns the component of the field $u$ at the dimension $dim$.

$$ u(x) = (u_0(x), u_1(x), \ldots, u_N(x)) \to u_{\text{dim}}(x) $$

Concrete implementations:

get_component(x: 'OnGrid', *, dim: 'int', params = None)

Slices the parameters of the field along the last dimensions, at the index specified by dim.

Parameters:

Returns:

shift_operator

Implements the shift operator $S(\Delta x)$ which is used to shift (spatially) a field $u$ by a constant $\Delta x$:

$$ v = S(\Delta x) u = u(x - \Delta x) $$

Concrete implementations:

shift_operator(x: 'Continuous', *, dx: 'object', params = None)

Shifts the field by dx using function composition.

Parameters:

Returns:

shift_operator(x: 'FiniteDifferences', *, dx = [0.0], params = None)

Shifts the field by dx using finite differences.

Parameters:

Returns:

shift_operator(x: 'FourierSeries', *, dx = [0], params = None)

Shifts the field by dx using the shift theorem in Fourier space.

Parameters:

Returns:

sum_over_dims

Reduces a vector field $u = (u_x, u_y, \dots)$ to a scalar field by summing over the dimensions:

$$ v = \sum_{i \in {x,y,\dots}} u_i $$

Concrete implementations:

sum_over_dims(x: 'Continuous', *, params = None)

sum_over_dims(x: 'OnGrid', *, params = None)

Utilities