Discretization API¶

jaxdf revolves around the concept of discretization.

We will call discretization family the mapping $\mathcal{D}$ that associates a function $f$ to a set of discrete parameters $\theta$

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

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

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

Example¶

A simple example of 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 have to provide the function $\mathcal{P}_N$, the parameters of the function and the domain where the function is defined

Let's look at the field over the domain using the on_grid method

To get the field at a specific location $x$, we can simply call the field with at the required coordinates.

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

Let's try to apply derivative operator to this newly defined field

The parameter of the field du_dx are the same as u, since for Continuous fields the gradient operator is evaluated using autograd, which is an operation that only affects the computational grah of a 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 use this knowledge, we first define a new discretization family from the Continuous one, and then we 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,get_field, however we can use knowledge of the formula for the get_fun.

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

To now 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 now we 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 the operators as methods of the fields, which are then redefined by the children classes. This was the approach used in jaxdf when it was first made public, and it still can be used

However, it can become cumbersome to deal with many derived methods for different kind of discretizations, especially if one starts to evaluate operators that accept multiple operands with different combinations of discretizations (e.g. dot products, +, heterogeneous differential operators, etc).

This problem elegantly resolved in some programming languages usign multiple-dispatch. One of the languages that notably supports multiple dispatch is the Julia language, and I suggest to look at the packages of the SciML echosystem if you are familiar with Julia and or interested in learning this language (those packages look great!).

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

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

The input of the fuction can be arbitrary types: if they are fields or any type which is traceable by jax, they will be traced. The function has a reserved, mandatory input keyword params, which is reserved for the parameters of the operator, like 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 use of the @operator decorator makes sense when the arguments are defined using type annotation. In that way, we are using the dispatch method of plum to define an implementation of that function which is specific for the annotated types.