A research framework for writing differentiable PDE discretizations in JAX

11/09/2021
by   Antonio Stanziola, et al.
0

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>

READ FULL TEXT
POST COMMENT

Comments

There are no comments yet.

Authors

page 6

page 7

10/04/2021

Improved architectures and training algorithms for deep operator networks

Operator learning techniques have recently emerged as a powerful tool fo...
05/21/2022

Spectral Neural Operators

A plentitude of applications in scientific computing requires the approx...
06/03/2020

dynoNet: a neural network architecture for learning dynamical systems

This paper introduces a network architecture, called dynoNet, utilizing ...
11/24/2020

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

ADCME is a novel computational framework to solve inverse problems invol...
12/23/2021

Integrating Material Selection with Design Optimization via Neural Networks

The engineering design process often entails optimizing the underlying g...
07/21/2018

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.