Quickstart¶

This notebook is designed to showcase the primary features of jaxdf. Throughout the notebook, we will utilize several libraries from the jax ecosystem, including optax and diffrax. Although these libraries are not necessary for using jaxdf, our goal is to demonstrate jaxdf's compatibility with generic jax-based libraries.

Let's begin by downloading an image that we will use in this notebook.

Heat Equation¶

For illustration, we'll simulate the heat equation.

$$ \frac{\partial}{\partial t}u = \nabla^2 u $$

NOTE: We don't enforce specific boundary conditions here as they are not yet implemented. At present, boundary conditions are implicitly defined by padding for convolutive operators in Finite Differences and Fourier as periodic BC. While this isn't ideal for properly integrating the heat equation, it works whenever we use some form of absorbing layer at the boundary, which is often the case in acoustics.

Obviously, a more suitable handling of generic boundary conditions would be a valuable addition to the package. Contributions are welcome 😊

We first define the domain where fields live. Subsequently, we specify the discretization family that will represent the inputs to the operator. In this case, we use Finite Differences, as given by the FiniteDifferences discretization.

Now, let's define the right-hand side of the heat equation. As the only operator required is the Laplacian, we can employ operators.laplacian to calculate it.

Integration using diffrax¶

In this section, we will integrate our heat equation using the diffrax library. diffrax offers a suite of differential equation solvers, making it ideal for our needs.

1. Setting up the PDE: First, we define our PDE field, which in this case is represented by the Laplacian of ( u ).

2. Setting the Solver and Time Steps:

   • We choose Tsit5 as our solver
   • We set our initial time t0 to 0 and the final time t_final to 20.
   • The time step dt is set to 0.1.
   • Additionally, with the SaveAt feature, we decide to save the results at specific, evenly spaced time points.

3. Solving the PDE: Using dfx.diffeqsolve, we pass in our term, the chosen solver, time configurations, and the initial field u to obtain the solution.

Let's dive into the code:

Anisotropic Diffusion example¶

Anisotropic diffusion is a more sophisticated form of heat diffusion where the diffusion rate can vary in different directions and magnitudes, often dependent on the underlying structure of the data. This is particularly useful in applications like image processing, where it can help preserve edges while diffusing noise.

Let's see how we can leverage jaxdf to define and solve an anisotropic diffusion problem:

The anisotropic diffusion equation can be represented as:

$$ \frac{\partial u}{\partial t} = \nabla \cdot (c(u) \nabla u) $$

Where:

• $u$ is the field of interest.
• $c(u)$ is the diffusion conductivity that varies depending on the gradient magnitude of ( u ).

To translate this to code:

• Divergence Operator: We start by defining the divergence of a vector field, divergence(u, stagger). This computes the rate at which density exits at each point, essential for understanding how diffusion spreads. Note that we are using staggered gradients, which are a feature of OnGrid fields.

• Diffusion Conductivity: conductivity_kernel(u) computes the conductivity based on the magnitude of the gradient of $u$. The kernel ensures that areas with high gradients (like edges) have lower conductivity, preserving features.

• Gradient Magnitude: norm(u) computes the magnitude of the gradient, which is essential to determine the diffusion rate.

• Complete Anisotropic Diffusion: anisotropic_diffusion(t, u, args) combines the above functions to compute the divergence of the conductivity-weighted gradient of $u$, yielding the right-hand side of our PDE.

By constructing the anisotropic diffusion operator using jaxdf, we've highlighted several features of the library:

• Composition: Ability to define complex differential operators, such as the divergence, gradient, and custom conductivity kernel, using other operators.
• Performance: Using @jit to just-in-time compile our functions, ensuring optimal performance during execution.

Customizing the Discretization¶

Exploring different discretizations is a powerful feature in jaxdf. For instance, if we wish to go from FiniteDifferences to a FourierSeries discretization, the transition is seamless with the following steps:

1. Define a New Field: Initialize your field using the desired FourierSeries discretization.
2. Invoke the Operator: Simply call the previously defined operator on the newly discretized field.

This flexibility ensures that researchers and developers can effortlessly experiment with various discretizations, often without restructuring their primary operators.

Automatic Differentiation with jaxdf¶

One of the significant strengths of the jaxdf framework is its seamless integration with jax's automatic differentiation capabilities. By leveraging jax's native transformations, users can efficiently compute gradients with respect to the arguments of operators.

Handling Operator Parameters¶

Operators in computational methods often come with associated parameters that play a critical role in their computation. In jaxdf, these parameters can be discretization-dependent, adapting based on the specific method in use.

For instance, the laplacian operator:

• In the context of FiniteDifferences, it's implemented using a stencil.
• For FourierSeries, it relies on transformations of the k-axis (spatial frequency axis).

The .default_params method provides insight into these parameters. When invoked with the same arguments as the operator, it returns the default parameters utilized.

By design, these parameters are statically compiled into the XLA function. However, jaxdf offers flexibility for scenarios such as when:

• Parameters have been manually adjusted.
• Parameters are sizeable, making static compilation inefficient.
• Different operators share the same parameters, as seen in Fourier methods. In such cases, users might desire to use a consistent variable across all.

To accommodate these scenarios, every operator in jaxdf has a params keyword. This reserved keyword allows users to override default parameters, giving them precise control over the computation.

jaxdf doesn't just stop at providing access to parameters. Users can also apply functional transformations directly to these parameters too!

To illustrate this, let's embark on an example. We'll optimize the stencil of the laplacian operator for a FiniteDifferences field to align its outcome more closely with the FourierSeries version.

More complex operators¶

We can of course construct more complex operators than the heat equation, using composition. Beyond the core functionality of operators, jaxdf grants direct access to the numerical parameters of the fields. This feature is invaluable for scenarios requiring fine-tuned control or experimentation with parameters.