Discretization API¶

jaxdf is built around the concept of discretization.

We define a discretization family as the mapping $\mathcal{D}$ that associates a function $f$ with a set of discrete parameters $\theta$:

$$ \theta \xrightarrow{\mathcal{D}}f $$

where $f \in \text{Range}(\mathcal{D})$, or equivalently, $f_\theta(x) = \mathcal{D}(\theta, x)$ is a function parametrized by $\theta$.

Here, $\theta$ is the discrete representation of $f$ over $\mathcal{D}$. This is analogous to the interpolation function defined in other libraries (see, for example, the Operator Discretization Library).

Example¶

A simple example of a discretization family is the set of $N$-th order polynomials on the interval $[0,1)$:

$$ \mathcal{P}_N(\theta,x) = \sum_{i=0}^N \theta_i x^i, \qquad \theta \in \mathbb{R}^{N+1} $$

In jaxdf, we construct such a field using the Continuous discretization. To do so, we provide the function $\mathcal{P}_N$, the parameters of the function, and the domain where the function is defined.

Let's visualize the field over the domain using the on_grid method:

To evaluate the field at a specific location $x$, we can simply call the field at the required coordinates:

Fields are pytrees based on equinox Module, so they natively work with jax.jit, jax.grad, etc.

Let's apply the derivative operator to this newly defined field:

The parameters of the field du_dx are the same as u, since for Continuous fields the gradient operator is evaluated using autograd, which affects only the computational graph of the function but not its inputs.

Customizing Discretizations¶

For a polynomial field, we actually know analytically how to compute derivatives:

$$ \frac{\partial}{\partial x}\mathcal{P}_N(\theta,x) = \frac{\partial}{\partial x} \sum_{i=0}^N \theta_i x^i = \sum_{i=0}^{N-1} i\theta_{i+1} x^i $$

To leverage this knowledge, we first define a new discretization family from the Continuous class, and then define the gradient method using the analytical formula.

This new class needs to initialize the parent class using the super().__init__() method. The input parameters are params, domain, and get_field. However, we can use our knowledge of the formula for the get_fun.

The params must be a PyTree compatible with jax.numpy (arrays, dictionaries of arrays, equinox modules, etc.). The domain attribute must be a jaxdf.geometry.Domain object defining the domain of the field. The get_field attribute must be a function that evaluates the field at a coordinate using the parameters contained in params, with the signature get_fun(params: Field.params, x: Union[jnp.ndarray, float]).

To define the derivative operator acting on polynomials, we have several options. One is to simply define a Python function that generates the Polynomial object resulting from the gradient computation.

Note that the code is fully differentiable and can be compiled:

However, note that we now have a derivative operator defined for all types, and we get incorrect results for fields that are not polynomials:

Operators and Multiple Dispatch¶

One way to avoid this is to implement operators as methods of the fields, which are then redefined by child classes. This was the approach used in jaxdf when it was first made public, and it can still be used.

However, it can become cumbersome to manage many derived methods for different discretizations, especially when evaluating operators that accept multiple operands with different combinations of discretizations (e.g., dot products, +, heterogeneous differential operators, etc.).

This problem is elegantly resolved in some programming languages using multiple dispatch. One language that notably supports multiple dispatch is Julia, and I recommend looking at the packages in the SciML ecosystem if you are familiar with Julia or interested in learning it (those packages are excellent!).

For those sticking with Python, jax, and jaxdf, we borrow these ideas using the Python multiple dispatch library plum. The jaxdf.operator decorator can be used to define new (parametric) operators:

@operator
def new_operator(x: Polynomial, y: Continuous, *, params=1.0):
    ... # Any jaxdf or jax-compatible code
    return Field(...)

The input of the function can be arbitrary types: if they are fields or any type traceable by jax, they will be traced. The function has a reserved, mandatory keyword argument params, which is used for the parameters of the operator, such as the coefficients of the stencil of a differential operator.

The output of the function can be any type, including jaxdf.Field or jax.numpy.ndarray.

The @operator decorator is most useful when arguments are defined using type annotations. This way, we leverage plum's dispatch method to define an implementation of that function specific to the annotated types.