diff --git a/partitioned-burgers-1d/.solver-nutils-dgalerkin/solver.py b/partitioned-burgers-1d/.solver-nutils-dgalerkin/solver.py new file mode 100755 index 000000000..28f1cfd20 --- /dev/null +++ b/partitioned-burgers-1d/.solver-nutils-dgalerkin/solver.py @@ -0,0 +1,148 @@ +# Original source: https://github.com/evalf/nutils/blob/d73749ff7d64c9ccafdbb88cd442f80b9448c118/examples/burgers.py + +from nutils import mesh, function, export, testing +from nutils.solver import System +from nutils.expression_v2 import Namespace +import treelog as log +import numpy as np +import itertools +import precice +import json +import os +import argparse + +def _generate_initial_condition(x_coords, ic_config, epoch): + np.random.seed(epoch) + ic_values = np.zeros(len(x_coords)) + if ic_config["type"] == "sinusoidal": + num_modes = ic_config.get("num_modes", 1) + superpositions = np.random.randint(2, num_modes + 1) + for _ in range(superpositions): + amp = np.random.uniform(0.1, 2) + k = np.random.randint(ic_config["wavenumber_range"][0], ic_config["wavenumber_range"][1] + 1) + phase_shift = np.random.uniform(0, 2 * np.pi) + ic_values += amp * np.sin(2 * np.pi * k * x_coords + phase_shift) + return ic_values + +def project_initial_condition(domain_min, domain_max, nelems, ic_config, epoch): + # 1. Generate a high-resolution "truth" on a fine grid + fine_res = nelems * 10 + fine_x = np.linspace(domain_min[0], domain_max[0], fine_res, endpoint=False) + fine_u = _generate_initial_condition(fine_x, ic_config, epoch) + + # 2. Average the high-resolution truth over each coarse cell + u_projected = np.zeros(nelems) + for i in range(nelems): + cell_start = i * 10 + cell_end = (i + 1) * 10 + u_projected[i] = np.mean(fine_u[cell_start:cell_end]) + + return u_projected + +def main(dim: int, + epoch: int = 0, + btype: str = 'discont', + degree: int = 1, + newtontol: float = 1e-5, + config_file: str = "precice-config.xml"): + + script_dir = os.path.dirname(os.path.abspath(__file__)) + with open(os.path.join(script_dir, "..", "python_participant", "config.json"), 'r') as f: + config = json.load(f)["solver"] + with open(os.path.join(script_dir, "ic_params.json"), 'r') as f: + ic_config = json.load(f)["initial_conditions"] + + config_path = os.path.join(script_dir, "..", f"{dim}d", config_file) + participant = precice.Participant("Solver", config_path, 0, 1) + + mesh_internal_name = f"Solver-Mesh-{dim}D-Internal" + mesh_boundaries_name = f"Solver-Mesh-{dim}D-Boundaries" + data_name = f"Data_{dim}D" + res = config[f"{dim}d_resolution"] + domain_min = config[f"{dim}d_domain_min"] + domain_max = config[f"{dim}d_domain_max"] + nelems = res[0] + + domain, geom = mesh.line(np.linspace(domain_min[0], domain_max[0], nelems + 1), periodic=True) + + # Define all nelems +1 nodes for evaluation + eval_coords_x = np.linspace(domain_min[0], domain_max[0], nelems + 1) + + # Define the nelems vertices for saving (all but the last) + trunc_coords_x = eval_coords_x[:-1] + internal_coords = np.array([trunc_coords_x, np.full(len(trunc_coords_x), domain_min[1])]).T + boundary_coords = np.array([[domain_min[0], domain_min[1]], [domain_max[0], domain_max[1]]]) + + internal_vertex_ids = participant.set_mesh_vertices(mesh_internal_name, internal_coords) + boundary_vertex_ids = participant.set_mesh_vertices(mesh_boundaries_name, boundary_coords) + + sample = domain.locate(geom, eval_coords_x, tol=1e-5) + + ns = Namespace() + ns.x = geom + ns.define_for('x', gradient='∇', normal='n', jacobians=('dV', 'dS')) + ns.u = domain.field('u', btype=btype, degree=degree) + ns.du = ns.u - function.replace_arguments(ns.u, 'u:u0') + ns.v = domain.field('v', btype=btype, degree=degree) + ns.t = function.field('t') + ns.dt = ns.t - function.field('t0') + ns.f = '.5 u^2' + ns.C = 1 + + res_pde = domain.integral('(v du / dt - ∇(v) f) dV' @ ns, degree=degree*2) + res_pde -= domain.interfaces.integral('[v] n ({f} - .5 C [u] n) dS' @ ns, degree=degree*2) + system = System(res_pde, trial='u', test='v') + + # Project the initial condition + u_averaged = project_initial_condition(domain_min, domain_max, nelems, ic_config, epoch) + ns.uic = domain.basis('discont', degree=0).dot(u_averaged) + + sqr = domain.integral('(u - uic)^2 dV' @ ns, degree=max(degree*2, 5)) + args = System(sqr, trial='u').solve() + + if participant.requires_initial_data(): + # Evaluate at all nodes + all_data = sample.eval(ns.u, arguments=args) + # Truncate last element + trunc_data = all_data[:-1] + boundary_data_values = np.array([trunc_data[0], trunc_data[-1]]) + + participant.write_data(mesh_internal_name, data_name, internal_vertex_ids, trunc_data) + participant.write_data(mesh_boundaries_name, data_name, boundary_vertex_ids, boundary_data_values) + + participant.initialize() + + args['t'] = 0. + + with log.iter.plain('timestep', itertools.count()) as steps: + for _ in steps: + if not participant.is_coupling_ongoing(): + break + + timestep = participant.get_max_time_step_size() + + args = system.step(timestep=timestep, arguments=args, timearg='t', suffix='0', tol=newtontol) + + all_data = sample.eval(ns.u, arguments=args) + trunc_data = all_data[:-1] + boundary_data_values = np.array([trunc_data[0], trunc_data[-1]]) + + participant.write_data(mesh_internal_name, data_name, internal_vertex_ids, trunc_data) + participant.write_data(mesh_boundaries_name, data_name, boundary_vertex_ids, boundary_data_values) + + participant.advance(timestep) + + participant.finalize() + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("dim", type=int, choices=[1, 2], help="Dimension of the simulation") + parser.add_argument('--config_file', type=str, default="precice-config.xml") + parser.add_argument('--btype', type=str, default='discont') + parser.add_argument('--degree', type=int, default=1) + parser.add_argument('--newtontol', type=float, default=1e-5) + parser.add_argument("--epoch", type=int, default=0, help="Current epoch number") + + args_cli = parser.parse_args() + + main(dim=args_cli.dim, epoch=args_cli.epoch, btype=args_cli.btype, degree=args_cli.degree, newtontol=args_cli.newtontol, config_file=args_cli.config_file) diff --git a/partitioned-burgers-1d/clean-tutorial.sh b/partitioned-burgers-1d/clean-tutorial.sh new file mode 120000 index 000000000..4713f5092 --- /dev/null +++ b/partitioned-burgers-1d/clean-tutorial.sh @@ -0,0 +1 @@ +../tools/clean-tutorial-base.sh \ No newline at end of file diff --git a/partitioned-burgers-1d/compare_burgers.py b/partitioned-burgers-1d/compare_burgers.py new file mode 100644 index 000000000..6fa0eeb9e --- /dev/null +++ b/partitioned-burgers-1d/compare_burgers.py @@ -0,0 +1,54 @@ +import os +import numpy as np +import matplotlib.pyplot as plt + +DATA_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__))) + +data_dir_dgalerkin = os.path.join(DATA_DIR, "solver-nutils-dgalerkin", "data-training") +data_dir_diffuse = os.path.join(DATA_DIR, "solver-scipy-fvolumes", "data-training") + +# file names are the same, folders are different +files_ = os.listdir(data_dir_dgalerkin) +files_ = sorted(f for f in files_ if f.endswith(".npz")) + + +file_num = 0 +timestep = 0 + +file_dgalerkin = os.path.join(data_dir_dgalerkin, files_[file_num]) +file_diffuse = os.path.join(data_dir_diffuse, files_[file_num]) + +data_dgalerkin = np.load(file_dgalerkin)["Solver-Mesh-1D-Internal"] +data_diffuse = np.load(file_diffuse)["Solver-Mesh-1D-Internal"] + +print(f"DG data shape: {data_dgalerkin.shape}") +print(f"FV data shape: {data_diffuse.shape}") + +plt.figure(figsize=(12, 5)) +plt.plot(data_dgalerkin[timestep, :], label='DG Solver', color='blue') +plt.plot(data_diffuse[timestep, :], label='FV Solver', color='orange', linestyle='--') +plt.title(f'Comparison at Timestep {timestep}') +plt.xlabel('Spatial Position') +plt.ylabel('Solution Value') +plt.legend() +plt.savefig(os.path.join(DATA_DIR, f'comparison_timestep_{timestep}.png')) + +# plot the imshow with unified colormap +vmin = min(data_dgalerkin.min(), data_diffuse.min()) +vmax = max(data_dgalerkin.max(), data_diffuse.max()) + +plt.figure(figsize=(12, 5)) +plt.subplot(1, 2, 1) +plt.imshow(data_dgalerkin.T, aspect='auto', cmap='viridis', vmin=vmin, vmax=vmax) +plt.title('DG Solver Evolution') +plt.ylabel('Spatial Position') +plt.xlabel('Timestep') +plt.colorbar(label='Solution Value') +plt.subplot(1, 2, 2) +plt.imshow(data_diffuse.T, aspect='auto', cmap='viridis', vmin=vmin, vmax=vmax) +plt.title('FV Solver Evolution') +plt.ylabel('Spatial Position') +plt.xlabel('Timestep') +plt.colorbar(label='Solution Value') +plt.tight_layout() +plt.show() diff --git a/partitioned-burgers-1d/comparison_DG_FV.png b/partitioned-burgers-1d/comparison_DG_FV.png new file mode 100644 index 000000000..f60abf4a8 Binary files /dev/null and b/partitioned-burgers-1d/comparison_DG_FV.png differ diff --git a/partitioned-burgers-1d/dirichlet-scipy/run.sh b/partitioned-burgers-1d/dirichlet-scipy/run.sh new file mode 100755 index 000000000..029a864e6 --- /dev/null +++ b/partitioned-burgers-1d/dirichlet-scipy/run.sh @@ -0,0 +1,4 @@ +#!/usr/bin/env bash +set -e -u + +python3 ../solver-scipy-fvolumes/solver.py Dirichlet diff --git a/partitioned-burgers-1d/generate_ic.py b/partitioned-burgers-1d/generate_ic.py new file mode 100644 index 000000000..f8cd79cf3 --- /dev/null +++ b/partitioned-burgers-1d/generate_ic.py @@ -0,0 +1,67 @@ +import numpy as np +import json +import os +import argparse +import matplotlib.pyplot as plt + +def _generate_initial_condition(x_coords, ic_config, epoch): + np.random.seed(epoch) + ic_values = np.zeros(len(x_coords)) + if ic_config["type"] == "sinusoidal": + num_modes = ic_config.get("num_modes", 1) + superpositions = np.random.randint(2, num_modes + 1) + for _ in range(superpositions): + amp = np.random.uniform(0.1, 2) + k = np.random.randint(ic_config["wavenumber_range"][0], ic_config["wavenumber_range"][1] + 1) + phase_shift = np.random.uniform(0, 2 * np.pi) + ic_values += amp * np.sin(2 * np.pi * k * x_coords + phase_shift) + return ic_values + +def project_initial_condition(domain_min, domain_max, nelems, ic_config, epoch): + # 1. Generate a high-resolution "truth" on a fine grid + fine_res = nelems * 10 + fine_x = np.linspace(domain_min, domain_max, fine_res, endpoint=False) + fine_u = _generate_initial_condition(fine_x, ic_config, epoch) + + # 2. Average the high-resolution truth over each coarse cell + u_projected = np.zeros(nelems) + for i in range(nelems): + cell_start = i * 10 + cell_end = (i + 1) * 10 + u_projected[i] = np.mean(fine_u[cell_start:cell_end]) + + return u_projected + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="Generate initial conditions for the Burgers equation simulation.") + parser.add_argument("--epoch", type=int, default=0, help="Seed for the random number generator to ensure reproducibility.") + args = parser.parse_args() + + script_dir = os.path.dirname(os.path.abspath(__file__)) + with open(os.path.join(script_dir, "ic_params.json"), 'r') as f: + config = json.load(f) + + ic_config = config["initial_conditions"] + domain_config = config["domain"] + + # full domain + full_domain_min = domain_config["full_domain_min"] + full_domain_max = domain_config["full_domain_max"] + nelems_total = domain_config["nelems_total"] + + # Generate IC + initial_condition = project_initial_condition(full_domain_min, full_domain_max, nelems_total, ic_config, args.epoch) + + output_path = os.path.join(script_dir, "initial_condition.npz") + np.savez(output_path, initial_condition=initial_condition) + + plt.figure(figsize=(8, 4)) + x_coords = np.linspace(full_domain_min, full_domain_max, nelems_total, endpoint=False) + plt.figure(figsize=(10, 5)) + plt.plot(x_coords, initial_condition, marker='.', linestyle='-') + plt.xlabel('Spatial Coordinate (x)') + plt.ylabel('Solution Value (u)') + plt.grid(True) + plt.savefig(os.path.join(script_dir, "initial_condition.png")) + print(f"Initial condition and plot saved to {output_path}") + plt.close() diff --git a/partitioned-burgers-1d/ic_params.json b/partitioned-burgers-1d/ic_params.json new file mode 100644 index 000000000..55b3fb0f8 --- /dev/null +++ b/partitioned-burgers-1d/ic_params.json @@ -0,0 +1,12 @@ +{ + "initial_conditions": { + "type": "sinusoidal", + "wavenumber_range": [1, 7], + "num_modes": 4 + }, + "domain": { + "nelems_total": 256, + "full_domain_min": 0.0, + "full_domain_max": 2.0 + } +} diff --git a/partitioned-burgers-1d/neumann-scipy/run.sh b/partitioned-burgers-1d/neumann-scipy/run.sh new file mode 100755 index 000000000..06b3dab56 --- /dev/null +++ b/partitioned-burgers-1d/neumann-scipy/run.sh @@ -0,0 +1,4 @@ +#!/usr/bin/env bash +set -e -u + +python3 ../solver-scipy-fvolumes/solver.py Neumann diff --git a/partitioned-burgers-1d/precice-config.xml b/partitioned-burgers-1d/precice-config.xml new file mode 100644 index 000000000..3cc66bc9a --- /dev/null +++ b/partitioned-burgers-1d/precice-config.xml @@ -0,0 +1,56 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/partitioned-burgers-1d/requirements.txt b/partitioned-burgers-1d/requirements.txt new file mode 100644 index 000000000..6f7bad844 --- /dev/null +++ b/partitioned-burgers-1d/requirements.txt @@ -0,0 +1,8 @@ +numpy +scipy +matplotlib +pyprecice + +# Optional: PyTorch for neural surrogate solver +# Uncomment if using surrogate-burgers/solver.py +# torch \ No newline at end of file diff --git a/partitioned-burgers-1d/run-partitioned-scipy.sh b/partitioned-burgers-1d/run-partitioned-scipy.sh new file mode 100755 index 000000000..0444bd945 --- /dev/null +++ b/partitioned-burgers-1d/run-partitioned-scipy.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +set -e -u + +rm precice-run -rf + +python3 generate_ic.py --epoch ${1:-0} + +cd dirichlet-scipy; pwd; ./run.sh & +cd ../neumann-scipy; pwd; ./run.sh && cd .. + +# full domain reference solution +cd solver-scipy-fvolumes; python3 solver.py None; cd .. + + +python3 visualize_partitioned_domain.py \ No newline at end of file diff --git a/partitioned-burgers-1d/solver-scipy-fvolumes/solver.py b/partitioned-burgers-1d/solver-scipy-fvolumes/solver.py new file mode 100644 index 000000000..dc265adb6 --- /dev/null +++ b/partitioned-burgers-1d/solver-scipy-fvolumes/solver.py @@ -0,0 +1,335 @@ +""" +This script solves the 1D viscous Burgers' equation for a partitioned domain problem. +The two participants: +- 'Dirichlet': Solves the left half of the domain. Receives Dirichlet BC from 'Neumann'. Provides Neumann BC to 'Neumann'. +- 'Neumann': Solves the right half of the domain. + +# |<---------------------Dirichlet---------------------->|<--------------------Neumann----------------------->| +# | du_dx=0 | ... | ... | ... | u[-1] <|> bc_right | bc_left <|> u[0] | ... | ... | ... | du_dx=0 | +# |<----------------------- u -------------------------->|<----------------------- u ------------------------>| +# | bc_left | | | bc_right | +""" +import numpy as np +from scipy.integrate import solve_ivp +from scipy.sparse import diags, identity +from scipy.optimize import root +import precice +import json +import os +import argparse + +def lax_friedrichs_flux(u_left, u_right): # Flux at cell interface between i-1/2 and i+1/2 + return 0.5 * (0.5 * u_left**2 + 0.5 * u_right**2) - 0.5 * (u_right - u_left) + +def flux_function(u_left, u_right): + return lax_friedrichs_flux(u_left, u_right) + +def burgers_rhs(t, u, dx, C, bc_left, bc_right): + # Ghost cells for BCs + u_padded = np.empty(len(u) + 2) + u_padded[0] = bc_left + u_padded[-1] = bc_right + u_padded[1:-1] = u + + flux = np.empty(len(u) + 1) + for i in range(len(flux)): + flux[i] = flux_function(u_padded[i], u_padded[i+1]) + + # viscosity + viscosity = C * (u_padded[2:] - 2 * u_padded[1:-1] + u_padded[:-2]) / dx**2 + return -(flux[1:] - flux[:-1]) / dx + viscosity + +def burgers_jacobian(t, u, dx, C, bc_left, bc_right): + n = len(u) + + u_padded = np.empty(n + 2) + u_padded[0] = bc_left + u_padded[-1] = bc_right + u_padded[1:-1] = u + + # --- viscosity (constant) --- + d_visc_di = -2 * C / dx**2 + d_visc_off = C / dx**2 + + # --- flux (Lax-Friedrichs) --- + df_du_di = -1 / dx + + # Upper diagonal + df_du_upper = -(0.5 * u_padded[2:n+1] - 0.5) / dx + + # Lower diagonal + df_du_lower = (0.5 * u_padded[0:n-1] + 0.5) / dx + + main_diag = df_du_di + d_visc_di + upper_diag = df_du_upper + d_visc_off + lower_diag = df_du_lower + d_visc_off + + jac = diags([main_diag, upper_diag, lower_diag], [0, 1, -1], shape=(n, n), format='csc') + + return jac + +def burgers_residual(u_new, u_old, dt, dx, C, bc_left, bc_right): + return u_new - u_old - dt * burgers_rhs(0, u_new, dx, C, bc_left, bc_right) + +def burgers_jacobian_residual(u_new, u_old, dt, dx, C, bc_left, bc_right): + n = len(u_new) + I = identity(n, format='csc') + J_rhs = burgers_jacobian(0, u_new, dx, C, bc_left, bc_right) + return (I - dt * J_rhs).toarray() # doesn't work with sparse matrix in root solver for some reason + +class BoundaryWrapper: + """ + Wrap the RHS and Jacobian to dynamically set BCs during the solve iterations with the updated state u. + """ + def __init__(self, dx, C, participant_name, u_from_neumann=None, du_dx_recv=None): + self.dx = dx + self.C = C + self.participant_name = participant_name + self.u_from_neumann = u_from_neumann + self.du_dx_recv = du_dx_recv + + def bc_left(self, u): + if self.participant_name == "Neumann": + return u[0] - self.du_dx_recv * self.dx + # zero gradient at outer boundary + elif self.participant_name == "Dirichlet": + return u[0] + else: + return u[0] + + def bc_right(self, u): + if self.participant_name == "Dirichlet": + return self.u_from_neumann + # zero gradient at outer boundary + elif self.participant_name == "Neumann": + return u[-1] + else: + return u[-1] + + def rhs(self, t, u): + bc_left = self.bc_left(u) + bc_right = self.bc_right(u) + return burgers_rhs(t, u, self.dx, self.C, bc_left, bc_right) + + def jac(self, t, u): + bc_left = self.bc_left(u) + bc_right = self.bc_right(u) + + J_rhs = burgers_jacobian(t, u, self.dx, self.C, bc_left, bc_right) + return J_rhs + + def rhs_residual(self, u_new, u_old, dt): + bc_left = self.bc_left(u_new) + bc_right = self.bc_right(u_new) + return burgers_residual(u_new, u_old, dt, self.dx, self.C, bc_left, bc_right) + + def jac_residual(self, u_new, u_old, dt): + bc_left = self.bc_left(u_new) + bc_right = self.bc_right(u_new) + return burgers_jacobian_residual(u_new, u_old, dt, self.dx, self.C, bc_left, bc_right) + +def main(participant_name: str): + script_dir = os.path.dirname(os.path.abspath(__file__)) + case_dir = os.path.abspath(os.path.join(script_dir, '..')) + run_dir = os.getcwd() + + config_path = os.path.join(case_dir, "precice-config.xml") + + if participant_name == 'None': + # read precice config to get t_final and dt + import re + print("Participant not specified. Running full domain without preCICE") + participant_name = None + + with open(config_path, 'r') as f: + precice_config = f.read() + max_time_match = re.search(r'', precice_config) + t_final = float(max_time_match.group(1)) + dt_match = re.search(r'', precice_config) + dt = float(dt_match.group(1)) + print(f"t_final = {t_final}, dt = {dt}") + else: + participant = precice.Participant(participant_name, config_path, 0, 1) + + # Read initial condition + with open(os.path.join(case_dir, "ic_params.json"), 'r') as f: + domain_config = json.load(f)["domain"] + + nelems_total = domain_config["nelems_total"] + nelems_local = nelems_total // 2 + full_domain_min = domain_config["full_domain_min"] + full_domain_max = domain_config["full_domain_max"] + dx = (full_domain_max - full_domain_min) / nelems_total + + # Set domain and preCICE setup + if participant_name == "Dirichlet": + mesh_name = "Dirichlet-Mesh" + read_data_name = "Velocity" + write_data_name = "Gradient" + local_domain_min = full_domain_min + local_domain_max = full_domain_min + nelems_local * dx + coupling_point = [[local_domain_max, 0.0]] + elif participant_name == "Neumann": + mesh_name = "Neumann-Mesh" + read_data_name = "Gradient" + write_data_name = "Velocity" + local_domain_min = full_domain_min + nelems_local * dx + local_domain_max = full_domain_max + coupling_point = [[local_domain_min, 0.0]] + else: #full domain run + local_domain_min = full_domain_min + local_domain_max = full_domain_max + nelems_local = nelems_total + + ic_data = np.load(os.path.join(case_dir, "initial_condition.npz")) + full_ic = ic_data['initial_condition'] + if participant_name == "Dirichlet": + u = full_ic[:nelems_local] + elif participant_name == "Neumann": + u = full_ic[nelems_local:] + else: + u = full_ic + + if participant_name is not None: + vertex_id = participant.set_mesh_vertices(mesh_name, coupling_point) + + if participant.requires_initial_data(): + if participant_name == "Dirichlet": + du_dx_send = (u[-1] - u[-2]) / dx # take forward difference inside domain for initial send + participant.write_data(mesh_name, write_data_name, vertex_id, [du_dx_send]) + if participant_name == "Neumann": + participant.write_data(mesh_name, write_data_name, vertex_id, [u[0]]) + + participant.initialize() + dt = participant.get_max_time_step_size() + + solution_history = {int(0): u.copy()} + + t = 0.0 + t_index = 0 + saved_t = 0.0 + C_viscosity = 1e-12 + aborted = False + + # --- Serial Coupling Loop --- + if participant_name == "Dirichlet": + while participant.is_coupling_ongoing(): + if participant.requires_writing_checkpoint(): + # print(f"[Dirichlet] Writing checkpoint at t={t:.4f}") + saved_u = u.copy() + saved_t = t + if participant.requires_reading_checkpoint(): + u = saved_u.copy() + t = saved_t + # print(f"[Dirichlet] Reading checkpoint at t={t:.4f}") + + u_from_neumann = participant.read_data(mesh_name, read_data_name, vertex_id, dt)[0] + + t_end = t + dt + wrapper = BoundaryWrapper(dx, C_viscosity, "Dirichlet", u_from_neumann=u_from_neumann) + + # sol = solve_ivp(wrapper.rhs, (t, t_end), u, method='BDF', t_eval=[t_end], jac=wrapper.jac) # BDF higher order implicit timestepping + # u = sol.y[:, -1] + # u = u + dt * burgers_rhs(t, u, dx, C_viscosity, wrapper.bc_left(u), wrapper.bc_right(u)) # Explicit Euler + + sol = root(wrapper.rhs_residual, u, args=(u, dt), jac=wrapper.jac_residual, method='hybr') + u = sol.x + + bc_right = wrapper.bc_right(u) + + du_dx_send = (bc_right - u[-1]) / dx + flux_across_interface = flux_function(u[-1], bc_right) + u_interface = (u[-1] + bc_right) / 2 + + participant.write_data(mesh_name, write_data_name, vertex_id, [du_dx_send]) + + print(f"[{participant_name:9s}] t={t:6.4f} | u_coupling={u_interface:8.4f} | du_dx={du_dx_send:8.4f} | flux_across={flux_across_interface:8.4f}") + + t = saved_t + dt + t_index = int(t/dt) + solution_history[t_index] = u.copy() + participant.advance(dt) + + elif participant_name == "Neumann": + while participant.is_coupling_ongoing(): + if participant.requires_writing_checkpoint(): + # print(f"[Neumann] Writing checkpoint at t={t:.4f}") + saved_u = u.copy() + saved_t = t + if participant.requires_reading_checkpoint(): + u = saved_u.copy() + t = saved_t + # print(f"[Neumann] Reading checkpoint at t={t:.4f}") + + du_dx_recv = participant.read_data(mesh_name, read_data_name, vertex_id, dt)[0] + + t_end = t + dt + wrapper = BoundaryWrapper(dx, C_viscosity, "Neumann", du_dx_recv=du_dx_recv) + # sol = solve_ivp(wrapper.rhs, (t, t_end), u, method='BDF', t_eval=[t_end], jac=wrapper.jac) # BDF higher order implicit timestepping + # u = sol.y[:, -1] + # u = u + dt * burgers_rhs(t, u, dx, C_viscosity, wrapper.bc_left(u), wrapper.bc_right(u)) # Explicit Euler + + sol = root(wrapper.rhs_residual, u, args=(u, dt), jac=wrapper.jac_residual, method='hybr') + u = sol.x + + bc_left = wrapper.bc_left(u) + flux_across_interface = flux_function(bc_left, u[0]) + du_dx = (u[0] - bc_left) / dx + u_interface = (bc_left + u[0]) / 2 + + participant.write_data(mesh_name, write_data_name, vertex_id, [u[0]]) + + print(f"[{participant_name:9s}] t={t:6.4f} | u_coupling={u_interface:8.4f} | du_dx={du_dx:8.4f} | flux_across={flux_across_interface:8.4f}") + + t = saved_t + dt + t_index = int(t/dt) + solution_history[t_index] = u.copy() + participant.advance(dt) + + if participant_name is not None: + # Finalize and save data to npz array + participant.finalize() + output_filename = os.path.join(run_dir, f"{participant_name.lower()}.npz") + else: + output_filename = os.path.join(script_dir, "full_domain.npz") + print("Starting standalone simulation without preCICE") + bc_left, bc_right = 0, 0 + + while t + dt < t_final: + + print(f"[Standalone ] t={t:6.4f}") + t_end = t + dt + wrapper = BoundaryWrapper(dx, C_viscosity, "None") + # sol = solve_ivp(wrapper.rhs, (t, t_end), u, method='BDF', t_eval=[t_end], jac=wrapper.jac) # BDF higher order implicit timestepping + # u = sol.y[:, -1] + # u = u + dt * burgers_rhs(t, u, dx, C_viscosity, wrapper.bc_left(u), wrapper.bc_right(u)) # Explicit Euler + + sol = root(wrapper.rhs_residual, u, args=(u, dt), jac=wrapper.jac_residual, method='hybr') + u = sol.x + + t = t + dt + t_index = int(t/dt) + solution_history[t_index] = u.copy() + + if not aborted: + + cell_centers_x = np.linspace(local_domain_min + dx/2, local_domain_max - dx/2, nelems_local) + internal_coords = np.array([cell_centers_x, np.zeros(nelems_local)]).T + + sorted_times_index = sorted(solution_history.keys()) + final_solution = np.array([solution_history[t_index] for t_index in sorted_times_index]) + + np.savez( + output_filename, + internal_coordinates=internal_coords, + **{"Solver-Mesh-1D-Internal": final_solution} + ) + print(f"Results saved to {output_filename}") + else: + raise RuntimeError("Simulation aborted.") + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("participant", help="Name of the participant", choices=['Dirichlet', 'Neumann', 'None']) + args = parser.parse_args() + main(args.participant) \ No newline at end of file diff --git a/partitioned-burgers-1d/surrogate-burgers/CNN_RES_UNROLL_5.pth b/partitioned-burgers-1d/surrogate-burgers/CNN_RES_UNROLL_5.pth new file mode 100644 index 000000000..bbbea3245 Binary files /dev/null and b/partitioned-burgers-1d/surrogate-burgers/CNN_RES_UNROLL_5.pth differ diff --git a/partitioned-burgers-1d/surrogate-burgers/run.sh b/partitioned-burgers-1d/surrogate-burgers/run.sh new file mode 100755 index 000000000..10372b843 --- /dev/null +++ b/partitioned-burgers-1d/surrogate-burgers/run.sh @@ -0,0 +1,3 @@ +#!/bin/bash +cd "$(dirname "$0")" +python3 solver.py \ No newline at end of file diff --git a/partitioned-burgers-1d/surrogate-burgers/solver.py b/partitioned-burgers-1d/surrogate-burgers/solver.py new file mode 100644 index 000000000..a5e9956bd --- /dev/null +++ b/partitioned-burgers-1d/surrogate-burgers/solver.py @@ -0,0 +1,135 @@ +import torch +import numpy as np +import precice +import os +import sys +import json +import argparse +import time + +script_dir = os.path.dirname(os.path.abspath(__file__)) +project_root = os.path.abspath(os.path.join(script_dir, '..', '..', '..')) +sys.path.append(project_root) + +from neural_surrogate.model import CNN_RES + +def main(): + participant_name = "Neumann" + + case_dir = os.path.abspath(os.path.join(script_dir, '..')) + config_path = os.path.join(case_dir, "precice-config.xml") + + participant = precice.Participant(participant_name, config_path, 0, 1) + + # Read initial condition + with open(os.path.join(case_dir, "ic_params.json"), 'r') as f: + domain_config = json.load(f)["domain"] + + nelems_total = domain_config["nelems_total"] + nelems_local = nelems_total // 2 + full_domain_min = domain_config["full_domain_min"] + full_domain_max = domain_config["full_domain_max"] + dx = (full_domain_max - full_domain_min) / nelems_total + + # Set domain and preCICE setup + mesh_name = "Neumann-Mesh" + read_data_name = "Flux" + write_data_name = "Velocity" + local_domain_min = full_domain_min + nelems_local * dx + coupling_point = [[local_domain_min, 0.0]] + vertex_id = participant.set_mesh_vertices(mesh_name, coupling_point) + + ic_data = np.load(os.path.join(case_dir, "initial_condition.npz")) + full_ic = ic_data['initial_condition'] + u = full_ic[nelems_local:] + solution_history = {0.0: u.copy()} + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model_path = os.path.join(script_dir, "CNN_RES_UNROLL_5.pth") + + NUM_RES_BLOCKS = 4 + KERNEL_SIZE = 5 + + INPUT_SIZE = 128 + 2 # +2 for ghost cells + HIDDEN_SIZE = 256 + OUTPUT_SIZE = 128 + + model = CNN_RES( + hidden_channels=HIDDEN_SIZE, + num_blocks=NUM_RES_BLOCKS, + kernel_size=KERNEL_SIZE + ) + + if not os.path.exists(model_path): + print(f"Model file not found at {model_path}") + sys.exit(1) + + + model.load_state_dict(torch.load(model_path, map_location=device)) + model.to(device) + model.eval() + + print("Neural surrogate model loaded successfully.") + + if participant.requires_initial_data(): + participant.write_data(mesh_name, write_data_name, vertex_id, [u[0]]) + + participant.initialize() + + dt = participant.get_max_time_step_size() + t = 0.0 + saved_t = 0.0 + + # Main Coupling Loop + with torch.no_grad(): + while participant.is_coupling_ongoing(): + if participant.requires_writing_checkpoint(): + saved_u = u.copy() + saved_t = t + if participant.requires_reading_checkpoint(): + u = saved_u.copy() + t = saved_t + + participant.write_data(mesh_name, write_data_name, vertex_id, [u[0]]) + + du_dx_bc = participant.read_data(mesh_name, read_data_name, vertex_id, dt)[0] + + # Calculate ghost cell value from received flux + bc_left = u[0] - dx * du_dx_bc + bc_right = 0.0 # Dirichlet on the far right wall + + u_padded = np.empty(len(u) + 2) + u_padded[0] = bc_left + u_padded[-1] = bc_right + u_padded[1:-1] = u + + input_tensor = torch.from_numpy(u_padded).float().unsqueeze(0).unsqueeze(0).to(device) + + output_tensor = model(input_tensor) + u = output_tensor.squeeze().cpu().numpy() + + t += dt + solution_history[t] = u.copy() + participant.advance(dt) + + # Finalize and save data to npz array + participant.finalize() + + run_dir = os.getcwd() + output_filename = os.path.join(run_dir, "surrogate.npz") + + cell_centers_x = np.linspace(local_domain_min + dx/2, local_domain_min + (nelems_local - 0.5) * dx, nelems_local) + internal_coords = np.array([cell_centers_x, np.zeros(nelems_local)]).T + + sorted_times = sorted(solution_history.keys()) + final_solution = np.array([solution_history[time] for time in sorted_times]) + + np.savez( + output_filename, + internal_coordinates=internal_coords, + **{"Solver-Mesh-1D-Internal": final_solution} + ) + print(f"[Surrogate] Results saved to {output_filename}") + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/partitioned-burgers-1d/visualize_partitioned_domain.py b/partitioned-burgers-1d/visualize_partitioned_domain.py new file mode 100644 index 000000000..d1876245c --- /dev/null +++ b/partitioned-burgers-1d/visualize_partitioned_domain.py @@ -0,0 +1,119 @@ +import numpy as np +import matplotlib.pyplot as plt +import os + +TIMESTEP_TO_PLOT = 50 #eg. 0, 1, ..., n, ... ,-1 + +CASE_DIR = os.path.dirname(os.path.abspath(__file__)) +DIRICHLET_DATA_PATH = os.path.join(CASE_DIR, "dirichlet-scipy", "dirichlet.npz") +NEUMANN_DATA_PATH = os.path.join(CASE_DIR, "neumann-scipy", "neumann.npz") +# NEUMANN_DATA_PATH = os.path.join(CASE_DIR, "surrogate-burgers", "surrogate.npz") + +GROUND_TRUTH_DATA_PATH = os.path.join(CASE_DIR, "solver-scipy-fvolumes", "full_domain.npz") +if os.path.exists(GROUND_TRUTH_DATA_PATH): + print(f"Found ground truth data at {GROUND_TRUTH_DATA_PATH}") + gt_exists = True +else: + print(f"Ground truth data not found at {GROUND_TRUTH_DATA_PATH}.\nPlease run python3 solver-scipy-fvolumes/solver.py None.") + gt_exists = False + +print(f"Loading data from {DIRICHLET_DATA_PATH}") +data_d = np.load(DIRICHLET_DATA_PATH) +coords_d = data_d['internal_coordinates'] +solution_d = data_d['Solver-Mesh-1D-Internal'] + +print(f"Loading data from {NEUMANN_DATA_PATH}") +data_n = np.load(NEUMANN_DATA_PATH) +coords_n = data_n['internal_coordinates'] +solution_n = data_n['Solver-Mesh-1D-Internal'] + +full_coords = np.concatenate((coords_d[:, 0], coords_n[:, 0])) +full_solution_history = np.concatenate((solution_d, solution_n), axis=1) + +print(f"Full domain shape: {full_solution_history.shape}") + +if gt_exists: + print(f"Loading ground truth data from {GROUND_TRUTH_DATA_PATH}") + data_gt = np.load(GROUND_TRUTH_DATA_PATH) + coords_gt = data_gt['internal_coordinates'] + solution_gt = data_gt['Solver-Mesh-1D-Internal'] + + print(f"Ground truth shape: {solution_gt.shape}") + +# --- plot single timestep --- +plt.figure(figsize=(10, 5), dpi=200) +plt.plot(full_coords, full_solution_history[TIMESTEP_TO_PLOT, :], marker='.', linestyle='-', label='Partitioned domain') +if gt_exists: + plt.plot(coords_gt[:, 0], solution_gt[TIMESTEP_TO_PLOT, :], marker='x', linestyle='--', alpha=0.5, c="red", label='Ground truth') +plt.title(f'Solution at Timestep {TIMESTEP_TO_PLOT}') +plt.xlabel('Spatial Coordinate (x)') +plt.ylabel('Solution Value (u)') +plt.grid(True) +plt.legend() +plt.savefig(os.path.join(CASE_DIR, f'full_domain_timestep_slice.png')) +print(f"Saved plot to full_domain_timestep_slice.png") + +if gt_exists: + # residual + residual = full_solution_history[TIMESTEP_TO_PLOT, :] - solution_gt[TIMESTEP_TO_PLOT, :] + l2_norm_residual = np.linalg.norm(residual) + l2_norm_gt = np.linalg.norm(solution_gt[TIMESTEP_TO_PLOT, :]) + relative_residual = l2_norm_residual / l2_norm_gt if l2_norm_gt > 1e-9 else 0.0 + + nelems_total = solution_gt.shape[1] + interface_idx = nelems_total // 2 - 1 + dx = coords_gt[1, 0] - coords_gt[0, 0] + + # t = 0 + u0_gt = solution_gt[0, :] + val_at_interface_t0 = (u0_gt[interface_idx] + u0_gt[interface_idx + 1]) / 2.0 + grad_at_interface_t0 = (u0_gt[interface_idx + 1] - u0_gt[interface_idx]) / dx + + # t = TIMESTEP_TO_PLOT + u_plot_gt = solution_gt[TIMESTEP_TO_PLOT, :] + val_at_interface_plot = (u_plot_gt[interface_idx] + u_plot_gt[interface_idx + 1]) / 2.0 + grad_at_interface_plot = (u_plot_gt[interface_idx + 1] - u_plot_gt[interface_idx]) / dx + + print("---") + print("Ground truth u at interface:") + print(f" t=0: u = {val_at_interface_t0:8.4f}, du/dx = {grad_at_interface_t0:8.4f}") + print(f" t={TIMESTEP_TO_PLOT}: u = {val_at_interface_plot:8.4f}, du/dx = {grad_at_interface_plot:8.4f}") + print() + print(f"Residual at t={TIMESTEP_TO_PLOT}:") + print(f" L2 Norm of Absolute Residual: {l2_norm_residual:10.6e}") + print(f" L2 Norm of Relative Residual: {relative_residual:10.6e}") + print("---") + +# --- plot gradient at single timestep --- +solution_slice = full_solution_history[TIMESTEP_TO_PLOT, :] +du_dx = np.gradient(solution_slice, full_coords) + +plt.figure(figsize=(10, 5), dpi=200) +plt.plot(full_coords, du_dx, marker='.', linestyle='-', label='Partitioned domain') + +if gt_exists: + solution_gt_slice = solution_gt[TIMESTEP_TO_PLOT, :] + du_dx_gt = np.gradient(solution_gt_slice, coords_gt[:, 0]) + plt.plot(coords_gt[:, 0], du_dx_gt, marker='x', linestyle='--', alpha=0.5, c="red", label='Ground truth') + +plt.title(f'Gradient (du/dx) at Timestep {TIMESTEP_TO_PLOT}') +plt.xlabel('Spatial Coordinate (x)') +plt.ylabel('Gradient Value (du/dx)') +plt.grid(True) +plt.legend() +plt.savefig(os.path.join(CASE_DIR, f'gradient_timestep_slice.png')) +print(f"Saved plot to gradient_timestep_slice.png") +plt.close() + +# --- plot time evolution --- +plt.figure(figsize=(10, 6)) +plt.imshow(full_solution_history.T, aspect='auto', cmap='viridis', origin='lower', + extent=[0, full_solution_history.shape[0], full_coords.min(), full_coords.max()]) +plt.colorbar(label='Solution Value (u)') +plt.title('Time Evolution of Partitioned Burgers eq.') +plt.xlabel('Timestep') +plt.ylabel('Spatial Coordinate (x)') +plt.tight_layout() +plt.savefig(os.path.join(CASE_DIR, 'full_domain_evolution.png')) +print("Saved plot to full_domain_evolution.png") +plt.close() \ No newline at end of file