Time discrete: advection of a 3D level-set functionΒΆ
Solves the advection of a 3D level-set function.
\[\partial_t u + a \cdot \nabla u = 0\]
where \(a: \mathbb{R}^3 \times (0, T) \to \mathbb{R}^3\) is a time-dependent velocity field and \(u: \mathbb{R}^3 \times (0, T) \to \mathbb{R}\) is the unknown function.
The initial condition is the level-set function of a sphere; the sphere is deformed by the velocity field, and goes back to the initial configuration at time \(T\).
The equation is solved using the neural semi-Lagrangian scheme with a natural gradient projector.
[1]:
# %%
import matplotlib.pyplot as plt
import numpy as np
import torch
from scimba_torch import PI
from scimba_torch.approximation_space.nn_space import NNxSpace
from scimba_torch.domain.meshless_domain.domain_3d import Cube3D
from scimba_torch.integration.monte_carlo import DomainSampler, TensorizedSampler
from scimba_torch.integration.monte_carlo_parameters import UniformParametricSampler
from scimba_torch.neural_nets.coordinates_based_nets.mlp import GenericMLP
from scimba_torch.numerical_solvers.temporal_pde.neural_semilagrangian import (
Characteristic,
NeuralSemiLagrangian,
)
from scimba_torch.numerical_solvers.temporal_pde.time_discrete import (
TimeDiscreteNaturalGradientProjector,
)
from scimba_torch.physical_models.temporal_pde.advection_diffusion_equation import (
AdvectionReactionDiffusionDirichletStrongForm,
)
from scimba_torch.utils.scimba_tensors import LabelTensor
DEFAULT_BOUNDS = [(0, 1), (0, 1), (0, 1)]
def f_ini(x, mu):
x1, x2, x3 = x.get_components()
x0 = 0.35
return (x1 - x0) ** 2 + (x2 - x0) ** 2 + (x3 - x0) ** 2 - 0.15**2
def a(t, x, mu):
x1, x2, x3 = x.get_components()
pi_x1 = PI * x1
pi_x2 = PI * x2
pi_x3 = PI * x3
cos_pi_t = torch.cos(PI * t / 3)
v1 = torch.sin(pi_x1) ** 2 * torch.sin(2 * pi_x2) * torch.sin(2 * pi_x3) * cos_pi_t
v2 = torch.sin(pi_x2) ** 2 * torch.sin(2 * pi_x1) * torch.sin(2 * pi_x3) * cos_pi_t
v3 = torch.sin(pi_x3) ** 2 * torch.sin(2 * pi_x1) * torch.sin(2 * pi_x2) * cos_pi_t
return torch.cat((2 * v1, -v2, -v3), dim=1)
def plot_implicit(fn, bbox=(0, 1), iter=None):
"""Create a plot of an implicit function.
Args:
fn: Implicit function (plot where fn==0).
bbox: The x, y, and z limits of plotted interval.
iter: Iteration number (optional).
"""
def make_xyz(x, y, z):
if not isinstance(x, np.ndarray):
x = x * torch.ones_like(torch.tensor(y).flatten().unsqueeze(1))
y = torch.tensor(y).flatten().unsqueeze(1)
z = torch.tensor(z).flatten().unsqueeze(1)
if not isinstance(y, np.ndarray):
y = y * torch.ones_like(torch.tensor(z).flatten().unsqueeze(1))
x = torch.tensor(x).flatten().unsqueeze(1)
z = torch.tensor(z).flatten().unsqueeze(1)
if not isinstance(z, np.ndarray):
z = z * torch.ones_like(torch.tensor(x).flatten().unsqueeze(1))
x = torch.tensor(x).flatten().unsqueeze(1)
y = torch.tensor(y).flatten().unsqueeze(1)
x_ = torch.cat((x, y, z), dim=1)
mu = torch.zeros((x_.shape[0], 0))
return LabelTensor(x_), LabelTensor(mu)
xmin, xmax, ymin, ymax, zmin, zmax = bbox * 3
fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
A = np.linspace(xmin, xmax, 100) # resolution of the contour
B = np.linspace(xmin, xmax, 15) # number of slices
A1, A2 = np.meshgrid(A, A) # grid on which the contour is plotted
for z in B: # plot contours in the XY plane
X, Y = A1, A2
Z = fn(*make_xyz(X, Y, z)).w[:, 0].detach().cpu().reshape((100, 100))
ax.contour(X, Y, Z + z, [z], zdir="z")
# [z] defines the only level to plot for this contour for this value of z
for y in B: # plot contours in the XZ plane
X, Z = A1, A2
Y = fn(*make_xyz(X, y, Z)).w[:, 0].detach().cpu().reshape((100, 100))
ax.contour(X, Y + y, Z, [y], zdir="y")
for x in B: # plot contours in the YZ plane
Y, Z = A1, A2
X = fn(*make_xyz(x, Y, Z)).w[:, 0].detach().cpu().reshape((100, 100))
ax.contour(X + x, Y, Z, [x], zdir="x")
# must set plot limits because the contour will likely extend
# way beyond the displayed level. Otherwise matplotlib extends the plot limits
# to encompass all values in the contour.
ax.set_zlim3d(zmin, zmax)
ax.set_xlim3d(xmin, xmax)
ax.set_ylim3d(ymin, ymax)
ax.set_aspect("equal")
plt.gca().set_rasterization_zorder(-1)
if iter:
plt.savefig(
f"{iter}_transport_3D_level_set_larger_truly_adaptive_large_time_step.pdf"
)
plt.show()
def plot_results(f, sampler, T=0, n_visu=85, iter=None):
plot_implicit(f)
# %%
class NonUniformSampler(DomainSampler):
def __init__(self, f, domain, **kwargs):
super().__init__(domain, **kwargs)
self.f = f
def sample_around_zeroes(self, n, sigma=0.05, keep=0.75):
candidates = LabelTensor(self.vol_sampler.sample(10 * n))
parameters = LabelTensor(torch.zeros((10 * n, 0)))
f_values = self.f(candidates, parameters).w[:, 0]
weights = torch.exp(-(f_values**2) / (2 * sigma**2))
return candidates[weights > keep].x
def sample_new_points(self, n: int):
"""Samples `n` new points from the volumetric domain and labels them by subdomains.
Args:
n (int): Number of points to sample.
Returns:
ScimbaTensor: A tensor of sampled points with their subdomain labels.
"""
pts_1 = self.sample_around_zeroes(n, sigma=0.002)
if pts_1.shape[0] >= n // 4:
pts_1 = pts_1[: n // 4]
pts_2 = self.sample_around_zeroes(n, sigma=0.01)
pts_2 = torch.cat((pts_1, pts_2), dim=0)
if pts_2.shape[0] >= n // 2:
pts_2 = pts_2[: n // 2]
pts_3 = self.sample_around_zeroes(n, sigma=0.05)
pts_3 = torch.cat((pts_2, pts_3), dim=0)
if pts_3.shape[0] >= 3 * n // 4:
pts_3 = pts_3[: 3 * n // 4]
remaining = n - pts_3.shape[0]
pts_unif = self.vol_sampler.sample(remaining)
pts = torch.cat((pts_3, pts_unif), dim=0)
labels = torch.zeros(n, dtype=int)
# Label points based on subdomains
for i, subdomain in enumerate(self.domain.list_subdomains):
condition = subdomain.is_inside(pts)
labels[condition] = i + 1
# Wrap the points and labels in a ScimbaTensor
data = LabelTensor(pts, labels)
data.x.requires_grad_()
return data
# %%
def solve_with_neural_sl(
T: float,
dt: float,
retrain_initial_condition=False,
N_c=10_000,
) -> NeuralSemiLagrangian:
torch.random.manual_seed(0)
domain_x = Cube3D(DEFAULT_BOUNDS, is_main_domain=True)
sampler_unif = TensorizedSampler(
[
DomainSampler(domain_x),
UniformParametricSampler([]),
]
)
space = NNxSpace(
1, 0, GenericMLP, domain_x, sampler_unif, layer_sizes=[35, 50, 50, 35]
)
pde = AdvectionReactionDiffusionDirichletStrongForm(
space, a=a, u0=f_ini, constant_advection=False, zero_diffusion=True
)
characteristic = Characteristic(pde, periodic=True)
sampler = TensorizedSampler(
[
NonUniformSampler(space.evaluate, domain_x),
# DomainSampler(domain_x),
UniformParametricSampler([]),
]
)
space.integrator = sampler
pde.space = space
characteristic = Characteristic(pde, periodic=True)
if retrain_initial_condition:
projector = TimeDiscreteNaturalGradientProjector(pde, rhs=f_ini)
scheme = NeuralSemiLagrangian(characteristic, projector)
scheme.initialization(epochs=250, n_collocation=N_c, verbose=True)
scheme.projector.save("ini_transport_3D_level_set_larger_truly_adaptive")
plot_results(space.evaluate, sampler.sample)
projector = TimeDiscreteNaturalGradientProjector(pde, rhs=f_ini)
projector.load("ini_transport_3D_level_set_larger_truly_adaptive")
projector.space.load_from_best_approx()
scheme = NeuralSemiLagrangian(characteristic, projector)
scheme.solve(
dt=dt,
final_time=T,
epochs=250,
n_collocation=N_c,
verbose=True,
# plot=plot_results,
save="transport_3D_level_set_larger_truly_adaptive_large_time_step",
additional_epochs_for_first_iteration=250,
)
plot_results(space.evaluate, sampler.sample)
return scheme, sampler
Uncomment the following to run the test (it can be long).
[2]:
#scheme, sampler = solve_with_neural_sl(
# 3,
# 3 / 10,
# N_c=64**3,
# retrain_initial_condition=True,
#)