Given a function \(\mathcal{P}: y → y2\), where \(y ∈ [0, 1]\), find a unknow function \(f\) that satisfies \(\mathcal{P}(f(x)) = x,\ ∀ x ∈ [0, 1]\).
A classical approach is to use a neural network to approximate the data points \(Ω\):
\[fθ(x) ≈ y,\ ∀ (x, y) ∈ Ω\]
where \(θ\) is the parameters of the neural network.
However, nowhere near the solution.
\[\mathcal{P}(fθ(x)) ≈ y,\ ∀ (x, y) ∈ Ω\]
Now, minimize the error between \(f2θ(x)\) and \(y\).
Neural networks that are trained to solve supervised learning tasks while respecting any given law of physics described by general nonlinear partial differential equations (PDE).
- Equation containing unknown functions and its partial derivative.
- Describe the relationship between independent variables, unknown functions and partial derivative.
- \(f(x, y) = ax + by + c\), where \(a, b, c\) are unknown parameters.
- \(u(x, y) = α u(x, y) + β f(x, y)\) where \(u\) is the unknown function.
- \(ux(x, y) = α uy(x, y) + β fxy(x, y)\) where \(ux\) is the partial derivative of \(u\) with respect to \(x\), \(uy\) is the partial derivative of \(u\) with respect to \(y\), and \(fxy\) is the partial derivative of \(f\) with respect to \(x\) and \(y\).
- \(\dot{u} = \frac{∂{u}}{∂{t}}\)
- \(uxy = {∂ 2u \over ∂ y\,∂ x}\)
- \(∇ u (x, y, z) = ux + uy + uz\)
- \(∇ ⋅ ∇ u(x, y, z) = Δ u(x,y,z) = uxx + uyy + uzz\)
- \(∇\) : nabla, or del.
\[Δ \varphi = 0\]
or
\[∇ ⋅ ∇ \varphi = 0\]
or, in a 3D space:
\[\frac {∂ 2f}{∂ x2}+\frac {∂ 2f}{∂ y2}+\frac {∂ 2f}{∂ z2} = 0\]
\[Δ \varphi = f\]
\[\dot{u} = \frac {∂ u}{∂ t} = α Δ u\]
where \(α\) is the thermal diffusivity.
\[\ddot {u}=c2∇ 2u\]
where \(c\) is the wave speed.
\[u_t + u u_x = ν uxx\]
- \(t\)
- temporal coordinate
- \(x\)
- spatial coordinate
- \(u(x, t)\)
- speed of fluid at the indicated spatial and temporal coordinates
- \(ν\)
- viscosity of fluid
For a equation \(∇2y+y=0\) in domain \(Ω\).
- Dirichlet boundary condition: \(y(x)=f(x)\quad ∀ x∈ ∂ Ω\)
- Neumann boundary condition: \(\frac {∂ y}{∂ \mathbf {n} }(\mathbf {x} )=f(\mathbf {x} )\quad ∀ \mathbf {x} ∈ ∂ Ω\)
- Where \(f\) is a known scalar function defined on the boundary domain \(∂ Ω\), \(\mathbf{n}\) denotes the (typically exterior) normal to the boundary.
- The normal derivative, which shows up on the left side, is defined as \(\frac {∂ y}{∂ \mathbf {n} }(\mathbf {x} )=∇ y(\mathbf {x} )⋅ \mathbf {\hat {n}} (\mathbf {x} )\), where \(\mathbf {\hat {n}}\) is the unit normal.
- Robin boundary condition
- Combine Dirichlet and Neumann boundary conditions.
- Periodic boundary condition
- Physics Informed Deep Learning (Part I): Data-driven Solutions of Nonlinear Partial Differential Equations[cite/ft/f:@raissiPhysicsInformedDeep2017]
- Physics Informed Deep Learning (Part II): Data-driven Discovery of Nonlinear Partial Differential Equations[cite/ft/f:@raissiPhysicsInformedDeep2017a]
- Data-driven solution and data-driven discovery
- Continuous time and discrete time models
General PDE Form:
\[u_t + \mathcal{N}[u] = 0,\ x ∈ Ω, \ t∈[0,T]\]
where:
- \(\mathcal{N}[u]\)
- nonlinear differential operator
- \(u(t, x)\)
- unknown function (solution).
- \(Ω\)
- spatial domain.
- \(t\)
- time.
- A neural network \(uθ ≈ u\), where \(θ\) is the parameters of the neural network.
- A physics informed neural network \(fθ = {uθ}t + \mathcal{N}[uθ]\).
- Target: \(fθ ≈ ut + \mathcal{N}[u]\) and \(uθ ≈ u\).
- \(\mathcal{L} = \mathcal{L}f + \mathcal{L}u\)
The equation:
\[u_t + u u_x = ν uxx\]
Here, already know \(ν = 0.01 / π\), \(x ∈ [-1, 1], t ∈ [0, 1]\),
Thus,
\[ut + uux - 0.01/π uxx = 0 \]
And the equation along with Dirichlet boundary conditions can be written as:
- \(u(0, x) = -sin(π x)\)
- \(u(t, -1) = u(t, 1) = 0\)
- Data:
- Boundary only data from boundary conditions.
- Input: \(\{t, x\}\)
- Output: \(u(t, x)\)
- Target: \(fθ ≈ ut + \mathcal{N}[u]\) and \(uθ ≈ u\).
- \(\mathcal{L} = \mathcal{L}f + \mathcal{L}u\)
def u_theta(theta, t, x):
# u_theta.apply(theta, t, x) to approx u(x, t)
return net(theta, t, x)
def f_theta(theta, t, x):
# See the auto diff cookbook
# https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html
u = u_theta.apply
u_t = jax.jacrev(u, argnums=1)(theta, t, x)
u_x = jax.jacrev(u, argnums=2)(theta, t, x)
u_xx = jax.hessian(u, argnums=2)(theta, t, x)
# or jax.jacfwd(jax.jacrev(u, argnums=2), argnums=2)
f = lambda: u_t + u * u_x - 0.01 * u_xx
return f- Train with MLPs with L-BFGS solver (quasi-newton method).
- Cannot use ReLU but tanh, because when we do the second order derivative, the ReLU will be 0.
General PDE Form:
\[u_t + \mathcal{N}[u;λ] = 0,\ x ∈ Ω, \ t∈[0,T]\]
where:
- \(\mathcal{N}[u;λ]\)
- nonlinear differential operator with parameters \(λ\).
- \(u(t, x)\)
- unknown function (solution).
- \(Ω\)
- spatial domain.
- \(t\)
- time.
The equations:
\[u_t + λ_1 (u u_x + v u_y) = -p_x + λ_2(uxx + uyy),\] \[v_t + λ_1 (u v_x + v v_y) = -p_y + λ_2(vxx + vyy)\],
where:
- \(u(t, x, y)\)
- \(x\)-component of the velocity field,
- \(v(t, x, y)\)
- \(y\)-component of the velocity field,
- \(p(t, x, y)\)
- pressure,
- \(λ\)
- the unknown parameters.
Additional physical constraints:
- Solutions to the Navier-Stokes equations are searched in the set of divergence-free functions, i.e.:
- \(ux + uy = 0\)
- which describes the conservation of mass of the fluid
- \(u\) and \(v\) can written as a latent function \(ψ(t, x, y)\) with an assumption:
- \(u = ψy, v = -ψx\)
- The neural network equations:
- \(f := u_t + λ_1 (u u_x + v u_y) + p_x - λ_2(uxx + uyy),\)
- \(g := v_t + λ_1 (u v_x + v v_y) + p_y - λ_2(vxx + vyy)\)
- Inptu: \(\{t,x,y,u,v\}\) with noisy.
- Output: \((ψ(t, x, y), p(t, x, y))\).
- Target:
- \(fθ ≈ f \)
- \(gθ ≈ g\)
- \(uθ ≈ u\)
- \(vθ ≈ v\)
JAX is Autograd and XLA, brought together for high-performance numerical computing and machine learning research. It provides composable transformations of Python+NumPy programs: differentiate, vectorize, parallelize, Just-In-Time compile to GPU/TPU, and more.
- \(f(x) = y\), always.
- non-pure function:
- IO operator:
print - No seed random function
time- Runtime error.
- IO operator:
- JAX (jax, jaxlib)
jaxjax.numpy
- Haiku (dm-haiku) from deepmind
- Modules
- Optax (optax) from deepmind
- Light
- Linear system optimizers (\(Ax = b\))
- JAXopt (jaxopt)
- Other optimizers.
- Jraph (jraph)
- Standardized data structures for graphs.
- JAX, M.D. (jax-md)
- JAX and Molecular Dynamics
- RLax (rlax), and Coax (coax)
- Reinforcement Learning
import jax
import jax.numpy as jnp
import haiku as hk
def _u(t, x):
return hk.MLP(jnp.concatenate([t, x], axis=-1), [10, 10, 1])
u = hk.transform_with_state(_u)fake_t = jnp.ones([batch, size])
fake_x = jnp.ones([batch, size])
# theta: params
# state: training state
# rng: random number generator
params, state = u.init(rng, fake_t, fake_x)
hk.experimental.tabulate(u)(fake_t, fake_x)def loss_fn(config, ...):
def _loss(params, t, x):
u_theta = u.apply(params, t, x)
...
loss = _f
return loss
return _loss
loss = loss_fn(config, ...)import optax
lr = optax.linear_schedule(
0.001, # init
0.001 / 10, # final
1, # steps change to final
150 # start linear decay after steps
)
opt = optax.adam(learning_rate=lr)
opt = optax.adamax(learning_rate=lr)import jaxopt
# Linear solver
solver = jaxopt.OptaxSolver(
loss,
opt,
maxiter=epochs,
...
)
# non-linear solver
solver = jaxopt.LBFGS(
loss,
maxiter=epochs,
...
)
opt_state = solver.init(params, state)
update = solver.update# init
params, state, opt_state, update
for batch in data:
params, state = update(params, state, batch)# Use pjit
from jax.experimental.maps import Mesh, ResourceEnv, thread_resources
from jax.experimental.pjit import PartitionSpec, pjit
mesh = Mesh(np.asarray(jax.devices(), dtype=object), ["data", ...])
thread_resources.env = ResourceEnv(physical_mesh=mesh, loops=())
update = pjit(
solver.update,
in_axis_resources=[
None, # params
None, # state
PartitionSpec("data"), # batch
],
out_axis_resources=None,
)- Find an inverse function of a parabola
- Classical
- Physics informed
- PDE
- PDE example
- PDE boundary
- PINNs
- Data-driven solution with continuous time
- Burgers’ equation
- Data-driven discovery with continuous time
- Navier-Stokes equation
- Data-driven solution with continuous time