A research framework for writing differentiable PDE discretizations in JAX

by   Antonio Stanziola, et al.

Differentiable simulators are an emerging concept with applications in several fields, from reinforcement learning to optimal control. Their distinguishing feature is the ability to calculate analytic gradients with respect to the input parameters. Like neural networks, which are constructed by composing several building blocks called layers, a simulation often requires computing the output of an operator that can itself be decomposed into elementary units chained together. While each layer of a neural network represents a specific discrete operation, the same operator can have multiple representations, depending on the discretization employed and the research question that needs to be addressed. Here, we propose a simple design pattern to construct a library of differentiable operators and discretizations, by representing operators as mappings between families of continuous functions, parametrized by finite vectors. We demonstrate the approach on an acoustic optimization problem, where the Helmholtz equation is discretized using Fourier spectral methods, and differentiability is demonstrated using gradient descent to optimize the speed of sound of an acoustic lens. The proposed framework is open-sourced and available at <https://github.com/ucl-bug/jaxdf>



There are no comments yet.


page 6

page 7


Improved architectures and training algorithms for deep operator networks

Operator learning techniques have recently emerged as a powerful tool fo...

Spectral Neural Operators

A plentitude of applications in scientific computing requires the approx...

dynoNet: a neural network architecture for learning dynamical systems

This paper introduces a network architecture, called dynoNet, utilizing ...

ADCME: Learning Spatially-varying Physical Fields using Deep Neural Networks

ADCME is a novel computational framework to solve inverse problems invol...

Integrating Material Selection with Design Optimization via Neural Networks

The engineering design process often entails optimizing the underlying g...

Decouple Learning for Parameterized Image Operators

Many different deep networks have been used to approximate, accelerate o...
This week in AI

Get the week's most popular data science and artificial intelligence research sent straight to your inbox every Saturday.