"""Postprocess the results of the 2D level-set transport problem.""" # %% import matplotlib.pyplot as plt import numpy as np import scipy.interpolate as spint import torch from contourpy import contour_generator from scimba_torch import PI from scimba_torch.approximation_space.nn_space import NNxSpace from scimba_torch.domain.meshless_domain.domain_2d import Square2D 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.collocation_projector import ( NaturalGradientProjector, ) from scimba_torch.numerical_solvers.temporal_pde.neural_semilagrangian import ( Characteristic, NeuralSemiLagrangian, ) 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)] def f_ini(x, mu): x1, x2 = x.get_components() return (x1 - 0.5) ** 2 + (x2 - 0.75) ** 2 - 0.15**2 def a(t, x, mu): x1, x2 = x.get_components() pi_x1 = PI * x1 pi_x2 = PI * x2 pi_t = PI * t / 8 v1 = torch.sin(pi_x1) ** 2 * torch.sin(2 * pi_x2) * torch.cos(pi_t) v2 = torch.sin(pi_x2) ** 2 * torch.sin(2 * pi_x1) * torch.cos(pi_t) return torch.cat((-v1, v2), dim=1) def get_prediction_on_grid(scheme, n_visu=768): f = scheme.pde.space.evaluate sampler = scheme.pde.space.integrator.sample x, mu = sampler(n_visu**2) x1 = torch.linspace(0, 1, n_visu) x2 = torch.linspace(0, 1, n_visu) x1, x2 = torch.meshgrid(x1, x2, indexing="ij") x = LabelTensor(torch.stack((x1.flatten(), x2.flatten()), dim=1)) x1, x2 = x.get_components() x1, x2 = x1.detach().cpu(), x2.detach().cpu() predictions = f(x, mu).w.detach().cpu() x1 = x1.reshape(n_visu, n_visu) x2 = x2.reshape(n_visu, n_visu) predictions = predictions.reshape(n_visu, n_visu) return x1, x2, predictions def plot_results(scheme, n_visu=768, iter=None): x1, x2, predictions = get_prediction_on_grid(scheme, n_visu) fig, ax = plt.subplots(1, 1, figsize=(6, 6), constrained_layout=True) predictions[predictions < 0] = -1 predictions[predictions > 0] = 1 ax.contourf(x1, x2, predictions, levels=0, cmap="turbo", zorder=-9) # fig.colorbar(contour, ax=ax, fraction=0.046, pad=0.04) ax.contour(x1, x2, predictions, levels=0, colors="white", linewidths=0.5) ax.set_title("Predictions") ax.set_aspect("equal") plt.gca().set_rasterization_zorder(-1) # plt.show() if iter: plt.savefig(f"{iter}_transport_2D_level_set.pdf") else: plt.show() # %% def load_and_postprocess_projector(file_name="final_level_set") -> NeuralSemiLagrangian: torch.random.manual_seed(0) domain_x = Square2D(DEFAULT_BOUNDS, is_main_domain=True) sampler = TensorizedSampler( [ DomainSampler(domain_x), UniformParametricSampler([]), ] ) if "larger" in file_name: layer_sizes = [35, 50, 35] elif "large" in file_name: layer_sizes = [25, 50, 25] elif "shallow" in file_name: layer_sizes = [40, 40] else: layer_sizes = [20, 20, 20] space = NNxSpace(1, 0, GenericMLP, domain_x, sampler, layer_sizes=layer_sizes) pde = AdvectionReactionDiffusionDirichletStrongForm( space, a=a, u0=f_ini, constant_advection=False, zero_diffusion=True ) characteristic = Characteristic(pde, periodic=True) projector = NaturalGradientProjector(space, rhs=f_ini) projector.load(file_name) scheme = NeuralSemiLagrangian(characteristic, projector) return scheme, sampler def compute_points(scheme): x1, x2, z = get_prediction_on_grid(scheme, n_visu=768) cont_gen = contour_generator(x=x1, y=x2, z=z) lines = cont_gen.lines(0) x = lines[0][:, 0] y = lines[0][:, 1] p = np.stack((x, y)) dp = p[:, 1:] - p[:, :-1] u = np.r_[0, np.cumsum((dp**2).sum(axis=0) ** 0.25)] spl = spint.make_interp_spline(u, p, axis=1) uu = np.linspace(u[0], u[-1], 1000) xx, yy = spl(uu) return x, y, xx, yy def plot_final_solution(scheme, save_to_csv=False): x, y, xx, yy = compute_points(scheme) plt.plot(xx, yy, label="approx", lw=0.5) t = np.linspace(0, 2 * np.pi, 1000) x = np.cos(t) * 0.15 + 0.5 y = np.sin(t) * 0.15 + 0.75 plt.plot(x, y, label="true", lw=0.5) plt.gca().set_aspect("equal") plt.legend() plt.show() if save_to_csv: np.savetxt( "circle_comparison.csv", np.stack((x, y, xx, yy)).T, delimiter=",", header="x_true,y_true,x_approx,y_approx", comments="", ) def plot_points(scheme, file_name): x, y, xx, yy = compute_points(scheme) plt.plot(xx, yy, lw=0.5) plt.gca().set_aspect("equal") plt.show() np.savetxt( f"{file_name}.csv", np.stack((xx, yy)).T, delimiter=",", header="x_approx,y_approx", comments="", ) def compute_error(scheme): x, y, xx, yy = compute_points(scheme) error_points = (x - 0.5) ** 2 + (y - 0.75) ** 2 - 0.15**2 error_spline = (xx - 0.5) ** 2 + (yy - 0.75) ** 2 - 0.15**2 print(f"error on points: {error_points.mean()}") print(f"error on spline: {error_spline.mean()}") return error_points, error_spline # %% if __name__ == "__main__": plot_other_nets = False plot_along_time = True if plot_other_nets: print("#" * 80) print("SHALLOW NETWORK") print("#" * 80) scheme, sampler = load_and_postprocess_projector("final_level_set_shallow") plot_final_solution(scheme) err_pts, err_spl = compute_error(scheme) print("#" * 80) print("SMALL NETWORK") print("#" * 80) scheme, sampler = load_and_postprocess_projector("final_level_set") plot_final_solution(scheme) err_pts, err_spl = compute_error(scheme) print("#" * 80) print("LARGE NETWORK") print("#" * 80) scheme, sampler = load_and_postprocess_projector("final_level_set_large") plot_final_solution(scheme) err_pts, err_spl = compute_error(scheme) print("#" * 80) print("LARGE NETWORK, ADAPTIVE SAMPLING") print("#" * 80) scheme, sampler = load_and_postprocess_projector( "final_level_set_large_adaptive" ) plot_final_solution(scheme) err_pts, err_spl = compute_error(scheme) print("#" * 80) print("LARGER NETWORK, ADAPTIVE SAMPLING") print("#" * 80) scheme, sampler = load_and_postprocess_projector("final_level_set_larger_adaptive") plot_final_solution(scheme, save_to_csv=True) err_pts, err_spl = compute_error(scheme) if plot_along_time: for i in [4, 8, 12, 16, 20, 30]: file_name = f"{i}_transport_2D_level_set_larger_adaptive" scheme, sampler = load_and_postprocess_projector(file_name) plot_points(scheme, file_name) # %%