r"""Solves the advection of a 3D level-set function. ..math:: \partial_t u + a \cdot \nabla u = 0 where :math:`a: \mathbb{R}^3 \times (0, T) \to \mathbb{R}^3` is a time-dependent velocity field and :math:`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 :math:`T`. The equation is solved using the neural semi-Lagrangian scheme with a natural gradient projector. """ # %% 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 # %% if __name__ == "__main__": scheme, sampler = solve_with_neural_sl( 3, 3 / 10, N_c=64**3, retrain_initial_condition=True, ) # %%