diff --git a/.gitignore b/.gitignore index 36b889d17..1ad71ee0a 100644 --- a/.gitignore +++ b/.gitignore @@ -103,3 +103,5 @@ share/ lib64 pyvenv.cfg + +examples/TwoFluidQuasiNeutralToy/runs/*/ \ No newline at end of file diff --git a/examples/TwoFluidQuasiNeutralToy/1D_Verification.py b/examples/TwoFluidQuasiNeutralToy/1D_Verification.py new file mode 100644 index 000000000..0e1df71c2 --- /dev/null +++ b/examples/TwoFluidQuasiNeutralToy/1D_Verification.py @@ -0,0 +1,267 @@ +from cunumpy import pi, cos, sin, zeros_like, ones_like +from struphy.io.options import EnvironmentOptions, BaseUnits, Time +from struphy.geometry import domains +from struphy.fields_background import equils +from struphy.topology import grids +from struphy.io.options import DerhamOptions +from struphy.initial import perturbations +from struphy.initial.base import GenericPerturbation +from struphy import Simulation +from struphy.linear_algebra.solver import SolverParameters + +import argparse +import os +import glob +import cunumpy as xp +import matplotlib.pyplot as plt + +from struphy.models.two_fluid_quasi_neutral_toy import TwoFluidQuasiNeutralToy + +parser = argparse.ArgumentParser() +parser.add_argument('bc', choices=['periodic', 'dirichlet_hom', 'dirichlet_inhom']) +args = parser.parse_args() +BC = args.bc + +name = f"runs/sim_1D_{BC}" + +env = EnvironmentOptions(sim_folder=name) + +B0 = 0 +nu = 10.0 +nu_e = 1.0 +Nel = (32, 1, 1) +p = (1, 1, 1) +epsilon = 1.0 +dt = 1 +Tend = 1 +sigma = 0 +tol = 1e-5 + +time_opts = Time(dt=dt, Tend=Tend) +domain = domains.Cuboid() +equil = equils.HomogenSlab(B0x=0, B0y=0, B0z=B0, beta=0, n0=0) +grid = grids.TensorProductGrid(num_elements=Nel) + +# ---- boundary conditions ---- +if BC == 'periodic': + derham_opts = DerhamOptions(degree=p, bcs=(None, None, None)) + +elif BC == 'dirichlet_hom': + derham_opts = DerhamOptions(degree=p, bcs=(("dirichlet", "dirichlet"), None, None)) + +elif BC == 'dirichlet_inhom': + derham_opts = DerhamOptions(degree=p, bcs=(("dirichlet", "dirichlet"), None, None)) + lifting_function_u = GenericPerturbation(lambda x, y, z: x + 1, comp=0, given_in_basis="physical") + lifting_function_ue = GenericPerturbation(lambda x, y, z: x, comp=0, given_in_basis="physical") + +# ---- manufactured solutions ---- +if BC == 'periodic': + def mms_phi(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return xp.sin(2 * xp.pi * x) + 1, xp.zeros_like(x), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + +elif BC == 'dirichlet_hom': + def mms_phi(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + +elif BC == 'dirichlet_inhom': + def mms_phi(x, y, z): + return xp.sin(2 * xp.pi * x), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return xp.sin(2 * xp.pi * x) + x + 1, xp.zeros_like(x), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return xp.sin(2 * xp.pi * x) + x, xp.zeros_like(x), xp.zeros_like(x) + +# ---- source terms ---- +if BC == 'periodic': + def source_function_u(x, y, z): + fx = 2.0 * pi * (cos(2 * pi * x) + 2 * nu * pi * sin(2 * pi * x)) + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + + def source_function_ue(x, y, z): + fx = -2.0 * pi * cos(2 * pi * x) + nu_e * 4.0 * pi**2 * sin(2 * pi * x) - sigma * sin(2 * pi * x) + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + +elif BC == 'dirichlet_hom': + def source_function_u(x, y, z): + fx = 2.0 * pi * (cos(2 * pi * x) + 2 * nu * pi * sin(2 * pi * x)) + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + + def source_function_ue(x, y, z): + fx = -2.0 * pi * cos(2 * pi * x) + nu_e * 4.0 * pi**2 * sin(2 * pi * x) - sigma * sin(2 * pi * x) + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + +elif BC == 'dirichlet_inhom': + def source_function_u(x, y, z): + fx = 2.0 * pi * (cos(2 * pi * x) + 2 * nu * pi * sin(2 * pi * x)) + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + + def source_function_ue(x, y, z): + fx = -2.0 * pi * cos(2 * pi * x) + (4.0 * nu_e * pi**2 - sigma) * sin(2 * pi * x) - sigma * x + fy = zeros_like(x) + fz = zeros_like(x) + return fx, fy, fz + +# ---- perturbation classes for MMS initial conditions ---- +class MMSIonVelocity(perturbations.Perturbation): + def __init__(self, comp=0): + self.comp = comp + self.given_in_basis = "physical" + + def __call__(self, x, y, z): + return mms_ion_u(x, y, z)[self.comp] + + +class MMSElectronVelocity(perturbations.Perturbation): + def __init__(self, comp=0): + self.comp = comp + self.given_in_basis = "physical" + + def __call__(self, x, y, z): + return mms_electron_u(x, y, z)[self.comp] + + +class MMSPotential(perturbations.Perturbation): + def __init__(self): + self.given_in_basis = "physical" + + def __call__(self, x, y, z): + return mms_phi(x, y, z)[0] + + +# ---- model ---- +model = TwoFluidQuasiNeutralToy() + +model.propagators.qn_full.options = model.propagators.qn_full.Options( + nu=nu, + nu_e=nu_e, + eps_norm=epsilon, + stab_sigma=sigma, + source_u=source_function_u, + source_ue=source_function_ue, + solver='gmres', + solver_params=SolverParameters(verbose=True, info=True, tol = tol), +) + +if BC == 'dirichlet_inhom': + model.ions.u.lifting_function = lifting_function_u + model.electrons.u.lifting_function = lifting_function_ue + +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +if __name__ == "__main__": + sim.run(verbose=True) + sim.pproc(verbose=True) + sim.load_plotting_data(verbose=True) + + simdata = sim.plotting_data + n1_vals = simdata.grids_log[0] + x = xp.linspace(0, 1, 100) + + os.makedirs(f'{name}/plots', exist_ok=True) + for f in glob.glob(f'{name}/plots/*.png'): + os.remove(f) + + def save_plot(n1_vals, numerical, analytical, ylabel, title, fname, t): + plt.plot(n1_vals, numerical, label='numerical') + plt.plot(x, analytical, '--', label='manufactured') + plt.plot(n1_vals, numerical, 'k.', markersize=4, label='n1 points') + plt.xlabel('x') + plt.ylabel(ylabel) + plt.title(f'{title} at t={t:.3f}') + plt.legend() + plt.grid(True) + plt.savefig(f'{name}/plots/{fname}_{t:.3f}.png', dpi=300) + plt.clf() + + for t in list(simdata.spline_values.ions.u_log.data.keys()): + u_ions = simdata.spline_values.ions.u_log.data[t] + u_electrons = simdata.spline_values.electrons.u_log.data[t] + phi = simdata.spline_values.em_fields.phi_log.data[t] + + mms_phi_x, _, _ = mms_phi(x, x*0, x*0) + mms_ion_ux, _, _ = mms_ion_u(x, x*0, x*0) + mms_el_ux, _, _ = mms_electron_u(x, x*0, x*0) + + save_plot(n1_vals, phi[0][:, 0, 0], mms_phi_x, 'φ', 'Potential φ', 'plot_potential', t) + save_plot(n1_vals, u_ions[0][:, 0, 0], mms_ion_ux, 'u_x', 'Ion velocity u_x', 'plot_ion_ux', t) + save_plot(n1_vals, u_electrons[0][:, 0, 0], mms_el_ux, 'u_x', 'Electron velocity', 'plot_electron_ux', t) + + # ---- lifting diagnostics ---- + if BC == 'dirichlet_inhom': + e1 = xp.linspace(0, 1, 200) + e2 = xp.array([0.5]) + e3 = xp.array([0.5]) + + for label, var in [('ion', model.ions.u), ('electron', model.electrons.u)]: + if var.spline_lift is None: + continue + + def _eval(fn, comp=0): + return fn(e1, e2, e3, squeeze_out=True)[comp] + + fig, axes = plt.subplots(1, 3, figsize=(12, 4)) + axes[0].plot(e1, _eval(var.spline_lift)); axes[0].set_title(f'{label}: spline_lift') + axes[1].plot(e1, _eval(var.spline_0)); axes[1].set_title(f'{label}: spline_0') + axes[2].plot(e1, _eval(var.boundary_spline)); axes[2].set_title(f'{label}: boundary_spline') + for ax in axes: + ax.set_xlabel('x'); ax.grid(True) + plt.tight_layout() + plt.savefig(f'{name}/plots/lifting_{label}.png', dpi=300) + plt.clf() + + # ---- source diagnostics ---- + prop = model.propagators.qn_full + e1 = xp.linspace(0, 1, 200) + e2 = xp.array([0.5]) + e3 = xp.array([0.5]) + zeros_e = xp.zeros_like(e1) + + for label, spline, src_fn, comp in [ + ('ion_source_x', prop._src_u, prop.options.source_u, 0), + ('electron_source_x', prop._src_ue, prop.options.source_ue, 0), + ]: + if spline is None: + print(f" {label}: None, skipping") + continue + vals_proj = spline(e1, e2, e3, squeeze_out=True)[comp] + vals_ref = src_fn(e1, zeros_e, zeros_e)[comp] + plt.figure(figsize=(8, 4)) + plt.plot(e1, vals_ref, '--', label='analytical') + plt.plot(e1, vals_proj, '-', label='projected (FE)') + plt.xlabel('x'); plt.title(f'{label}'); plt.legend(); plt.grid(True) + plt.savefig(f'{name}/plots/source_{label}.png', dpi=300) + plt.close() + print(f" -> saved {name}/plots/source_{label}.png") \ No newline at end of file diff --git a/examples/TwoFluidQuasiNeutralToy/2D_Verification.py b/examples/TwoFluidQuasiNeutralToy/2D_Verification.py new file mode 100644 index 000000000..85b671092 --- /dev/null +++ b/examples/TwoFluidQuasiNeutralToy/2D_Verification.py @@ -0,0 +1,350 @@ +from cunumpy import pi, cos, sin, zeros_like, ones_like +from struphy.io.options import EnvironmentOptions, BaseUnits, Time +from struphy.geometry import domains +from struphy.fields_background import equils +from struphy.topology import grids +from struphy.io.options import DerhamOptions +from struphy.initial import perturbations +from struphy.initial.base import GenericPerturbation +from struphy import Simulation +from struphy.linear_algebra.solver import SolverParameters + +import argparse +import os +import glob +import cunumpy as xp +import matplotlib.pyplot as plt + +from mpi4py import MPI + +from struphy.models.two_fluid_quasi_neutral_toy import TwoFluidQuasiNeutralToy + +# ------------------ args ------------------ +parser = argparse.ArgumentParser() +parser.add_argument('bc', choices=['periodic', 'dirichlet_inhom', 'neumann']) +args = parser.parse_args() +BC = args.bc + +name = f"runs/sim_2D_{BC}" + +# ------------------ setup ------------------ +env = EnvironmentOptions(sim_folder=name) + +B0 = 1 +nu = 10.0 +nu_e = 1.0 +Nel = (8, 8, 1) +p = (2, 2, 1) +epsilon = 1.0 +dt = 1 +Tend = 1 +sigma = 1 +tol = 1e-5 + +time_opts = Time(dt=dt, Tend=Tend) +domain = domains.Cuboid() +equil = equils.HomogenSlab(B0x=0, B0y=0, B0z=B0, beta=0, n0=0) +grid = grids.TensorProductGrid(num_elements=Nel) + +# ------------------ boundary conditions ------------------ +if BC == 'periodic': + derham_opts = DerhamOptions(degree=p, bcs=(None, None, None)) + +elif BC == 'neumann': + derham_opts = DerhamOptions( + degree=p, + bcs=(None, + ("free", "free"), + None) + ) + +elif BC == 'dirichlet_inhom': + derham_opts = DerhamOptions( + degree=p, + bcs=(("dirichlet", "dirichlet"), + ("dirichlet", "dirichlet"), + None) + ) + + lifting_function_u = [ + GenericPerturbation(lambda x, y, z: - x * y, comp=0, given_in_basis="physical"), + GenericPerturbation(lambda x, y, z: -xp.cos(2*pi*x)*xp.cos(2*pi*y) - x*y, comp=1, given_in_basis="physical"), + ] + lifting_function_ue = [ + GenericPerturbation(lambda x, y, z: - x * y - 1, comp=0, given_in_basis="physical"), + GenericPerturbation(lambda x, y, z: -xp.cos(4*pi*x)*xp.cos(4*pi*y) - x*y - 1, comp=1, given_in_basis="physical"), + ] + +# ------------------ manufactured solutions ------------------ +if BC == 'periodic': + def mms_phi(x, y, z): + return xp.cos(2*pi*x) + xp.sin(2*pi*y), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return -xp.sin(2*pi*x)*xp.sin(2*pi*y), -xp.cos(2*pi*x)*xp.cos(2*pi*y), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return -xp.sin(4*pi*x)*xp.sin(4*pi*y), -xp.cos(4*pi*x)*xp.cos(4*pi*y), xp.zeros_like(x) + +elif BC == 'neumann': + def mms_phi(x, y, z): + return xp.cos(2*pi*x) + xp.sin(2*pi*y), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return -xp.sin(2*pi*x)*xp.sin(2*pi*y), -xp.cos(2*pi*y), xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return -xp.sin(2*pi*x)*xp.sin(2*pi*y), -xp.cos(2*pi*y), xp.zeros_like(x) + +elif BC == 'dirichlet_inhom': + def mms_phi(x, y, z): + return xp.cos(2*pi*x) + xp.sin(2*pi*y), xp.zeros_like(x), xp.zeros_like(x) + + def mms_ion_u(x, y, z): + return -xp.sin(2*pi*x)*xp.sin(2*pi*y) - x*y, -xp.cos(2*pi*x)*xp.cos(2*pi*y) - x*y, xp.zeros_like(x) + + def mms_electron_u(x, y, z): + return -xp.sin(4*pi*x)*xp.sin(4*pi*y) - x*y - 1, -xp.cos(4*pi*x)*xp.cos(4*pi*y) - x*y - 1, xp.zeros_like(x) + +# ------------------ source terms ------------------ +if BC == 'periodic': + def source_function_u(x, y, z): + fx = -2*pi*xp.sin(2*pi*x) + B0/epsilon*xp.cos(2*pi*x)*xp.cos(2*pi*y) - nu*8*pi**2*xp.sin(2*pi*x)*xp.sin(2*pi*y) + fy = 2*pi*xp.cos(2*pi*y) - B0/epsilon*xp.sin(2*pi*x)*xp.sin(2*pi*y) - nu*8*pi**2*xp.cos(2*pi*x)*xp.cos(2*pi*y) + return fx, fy, zeros_like(x) + + def source_function_ue(x, y, z): + fx = 2*pi*xp.sin(2*pi*x) - B0/epsilon*xp.cos(4*pi*x)*xp.cos(4*pi*y) - nu_e*32*pi**2*xp.sin(4*pi*x)*xp.sin(4*pi*y) + sigma*xp.sin(4*pi*x)*xp.sin(4*pi*y) + fy = -2*pi*xp.cos(2*pi*y) + B0/epsilon*xp.sin(4*pi*x)*xp.sin(4*pi*y) - nu_e*32*pi**2*xp.cos(4*pi*x)*xp.cos(4*pi*y) + sigma*xp.cos(4*pi*x)*xp.cos(4*pi*y) + return fx, fy, zeros_like(x) + +elif BC == 'neumann': + def source_function_u(x, y, z): + fx = B0/epsilon*xp.cos(2*pi*y) - 2*pi*xp.sin(2*pi*x) - nu*8*pi**2*xp.sin(2*pi*x)*xp.sin(2*pi*y) + fy = 2*pi*xp.cos(2*pi*y) - nu*4*pi**2*xp.cos(2*pi*y) - B0/epsilon*xp.sin(2*pi*x)*xp.sin(2*pi*y) + return fx, fy, zeros_like(x) + + def source_function_ue(x, y, z): + fx = -B0/epsilon*xp.cos(2*pi*y) + 2*pi*xp.sin(2*pi*x) - nu_e*8*pi**2*xp.sin(2*pi*x)*xp.sin(2*pi*y) + sigma*xp.sin(2*pi*x)*xp.sin(2*pi*y) + fy = -2*pi*xp.cos(2*pi*y) - nu_e*4*pi**2*xp.cos(2*pi*y) + sigma*xp.cos(2*pi*y) + B0/epsilon*xp.sin(2*pi*x)*xp.sin(2*pi*y) + return fx, fy, zeros_like(x) + + +elif BC == 'dirichlet_inhom': + def source_function_u(x, y, z): + fx = -2*pi*xp.sin(2*pi*x) + B0/epsilon*xp.cos(2*pi*x)*xp.cos(2*pi*y) - nu*8*pi**2*xp.sin(2*pi*x)*xp.sin(2*pi*y) + fy = 2*pi*xp.cos(2*pi*y) - B0/epsilon*xp.sin(2*pi*x)*xp.sin(2*pi*y) - nu*8*pi**2*xp.cos(2*pi*x)*xp.cos(2*pi*y) + fx += - B0 / epsilon * x * y + fy += B0 * x * y / epsilon + return fx, fy, zeros_like(x) + + def source_function_ue(x, y, z): + fx = 2*pi*xp.sin(2*pi*x) - B0/epsilon*xp.cos(4*pi*x)*xp.cos(4*pi*y) - nu_e*32*pi**2*xp.sin(4*pi*x)*xp.sin(4*pi*y) + sigma*xp.sin(4*pi*x)*xp.sin(4*pi*y) + fy = -2*pi*xp.cos(2*pi*y) + B0/epsilon*xp.sin(4*pi*x)*xp.sin(4*pi*y) - nu_e*32*pi**2*xp.cos(4*pi*x)*xp.cos(4*pi*y) + sigma*xp.cos(4*pi*x)*xp.cos(4*pi*y) + fx += B0 / epsilon * (x * y + 1) - sigma * (x*y + 1) + fy += - B0 / epsilon * (x * y + 1) - sigma * (x*y + 1) + return fx, fy, zeros_like(x) + + + +class MMSIonVelocity(perturbations.Perturbation): + def __init__(self, comp=0): + self.comp = comp + self.given_in_basis = "physical" + + def __call__(self, x, y, z): + return mms_ion_u(x, y, z)[self.comp] + + +class MMSElectronVelocity(perturbations.Perturbation): + def __init__(self, comp=0): + self.comp = comp + self.given_in_basis = "physical" + + def __call__(self, x, y, z): + return mms_electron_u(x, y, z)[self.comp] + + +class MMSPotential(perturbations.Perturbation): + def __init__(self): + self.given_in_basis = "physical" + + def __call__(self, x, y, z): + return mms_phi(x, y, z)[0] + +# ------------------ model ------------------ +model = TwoFluidQuasiNeutralToy() + +model.propagators.qn_full.options = model.propagators.qn_full.Options( + nu=nu, + nu_e=nu_e, + eps_norm=epsilon, + stab_sigma=sigma, + source_u=source_function_u, + source_ue=source_function_ue, + solver='gmres', + solver_params=SolverParameters(verbose=True, info=True, tol=tol), +) + +if BC == 'dirichlet_inhom': + model.ions.u.lifting_function = lifting_function_u + model.electrons.u.lifting_function = lifting_function_ue + + # model.ions.u.add_perturbation(MMSIonVelocity(comp=0)) + # model.ions.u.add_perturbation(MMSIonVelocity(comp=1)) + # model.electrons.u.add_perturbation(MMSElectronVelocity(comp=0)) + # model.electrons.u.add_perturbation(MMSElectronVelocity(comp=1)) + # model.em_fields.phi.add_perturbation(MMSPotential()) + +# ------------------ simulation ------------------ +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------ run ------------------ +if __name__ == "__main__": + sim.run(verbose=True) + + if MPI.COMM_WORLD.Get_rank() == 0: + + sim.pproc(verbose=True) + sim.load_plotting_data(verbose=True) + + simdata = sim.plotting_data + + n1_vals = simdata.grids_log[0] + n2_vals = simdata.grids_log[1] + X, Y = xp.meshgrid(n1_vals, n2_vals, indexing='ij') + + x = xp.linspace(0, 1, 100) + y = xp.linspace(0, 1, 100) + Xf, Yf = xp.meshgrid(x, y, indexing='ij') + + os.makedirs(f'{name}/plots', exist_ok=True) + for f in glob.glob(f'{name}/plots/*.png'): + os.remove(f) + + def save_plot(numerical, analytical, title, fname, t): + fig, axes = plt.subplots(1, 2, figsize=(10, 4)) + im0 = axes[0].contourf(X, Y, numerical, levels=50) + axes[0].set_title('numerical') + plt.colorbar(im0, ax=axes[0]) + im1 = axes[1].contourf(Xf, Yf, analytical, levels=50) + axes[1].set_title('manufactured') + plt.colorbar(im1, ax=axes[1]) + fig.suptitle(f'{title} at t={t:.3f}') + plt.savefig(f'{name}/plots/{fname}_{t:.3f}.png', dpi=300) + plt.close(fig) + + for t in simdata.spline_values.ions.u_log.data.keys(): + u_ions = simdata.spline_values.ions.u_log.data[t] + u_electrons = simdata.spline_values.electrons.u_log.data[t] + phi = simdata.spline_values.em_fields.phi_log.data[t] + + mms_phi_x, _, _ = mms_phi(Xf, Yf, 0*Xf) + mms_ion_ux, mms_ion_uy, _ = mms_ion_u(Xf, Yf, 0*Xf) + mms_el_ux, mms_el_uy, _ = mms_electron_u(Xf, Yf, 0*Xf) + + save_plot(phi[0][:, :, 0], mms_phi_x, 'φ', 'plot_phi', t) + save_plot(u_ions[0][:, :, 0], mms_ion_ux, 'u_ix', 'plot_uix', t) + save_plot(u_ions[1][:, :, 0], mms_ion_uy, 'u_iy', 'plot_uiy', t) + save_plot(u_electrons[0][:, :, 0], mms_el_ux, 'u_ex', 'plot_uex', t) + save_plot(u_electrons[1][:, :, 0], mms_el_uy, 'u_ey', 'plot_uey', t) + + # ---- lifting diagnostics (dirichlet_inhom only) ---- + if BC == 'dirichlet_inhom': + e1 = xp.linspace(0, 1, 80) + e2 = xp.linspace(0, 1, 80) + e3 = xp.array([0.5]) + E1, E2 = xp.meshgrid(e1, e2, indexing='ij') + + for label, var, comp in [('ion_ux', model.ions.u, 0), + ('ion_uy', model.ions.u, 1), + ('electron_ux', model.electrons.u, 0), + ('electron_uy', model.electrons.u, 1)]: + if var.spline_lift is None: + print(f" {label}: spline_lift is None, skipping") + continue + + def _eval(fn): + return fn(e1, e2, e3, squeeze_out=True)[comp] + + fig, axes = plt.subplots(1, 3, figsize=(15, 4)) + for ax, fn, ttl in zip(axes, + [var.spline_lift, var.spline_0, var.boundary_spline], + ['lifting', 'zero-BC part', 'boundary spline']): + im = ax.contourf(E1, E2, _eval(fn), levels=50) + ax.set_title(f'{label}: {ttl}') + plt.colorbar(im, ax=ax) + out = f'{name}/plots/lifting_{label}.png' + plt.savefig(out, dpi=300) + plt.close(fig) + print(f" -> saved {out}") + + # ---- source diagnostics ---- + prop = model.propagators.qn_full + e1 = xp.linspace(0, 1, 80) + e2 = xp.linspace(0, 1, 80) + e3 = xp.array([0.5]) + E1, E2 = xp.meshgrid(e1, e2, indexing='ij') + zeros_E = xp.zeros_like(E1) + + for label, spline, src_fn, comp in [ + ('ion_source_x', prop._src_u, prop.options.source_u, 0), + ('ion_source_y', prop._src_u, prop.options.source_u, 1), + ('electron_source_x', prop._src_ue, prop.options.source_ue, 0), + ('electron_source_y', prop._src_ue, prop.options.source_ue, 1), + ]: + if spline is None: + print(f" {label}: None, skipping") + continue + + vals_proj = spline(e1, e2, e3, squeeze_out=True)[comp] + vals_ref = src_fn(E1, E2, zeros_E)[comp] + + fig, axes = plt.subplots(1, 2, figsize=(10, 4)) + im0 = axes[0].contourf(E1, E2, vals_proj, levels=50) + axes[0].set_title('projected (FE)') + plt.colorbar(im0, ax=axes[0]) + im1 = axes[1].contourf(E1, E2, vals_ref, levels=50) + axes[1].set_title('reference (analytical)') + plt.colorbar(im1, ax=axes[1]) + fig.suptitle(label) + out = f'{name}/plots/source_{label}.png' + plt.savefig(out, dpi=300) + plt.close(fig) + print(f" -> saved {out}") + + if BC == 'dirichlet_inhom': + y_check = xp.linspace(0, 1, 80) + x_check = xp.linspace(0, 1, 80) + z_check = xp.array([0.5]) + + # comp=0: normal trace at x=0 and x=1 + for x_bnd, label in [(0.0, 'x=0'), (1.0, 'x=1')]: + x_bnd_arr = xp.array([x_bnd]) + mms_vals = mms_ion_u(x_bnd_arr, y_check, z_check)[0] + lift_vals = model.ions.u.spline_lift(x_bnd_arr, y_check, z_check, squeeze_out=True)[0] + print(f"ion ux normal trace diff at {label}: max={xp.max(xp.abs(mms_vals - lift_vals)):.3e}") + + mms_vals = mms_electron_u(x_bnd_arr, y_check, z_check)[0] + lift_vals = model.electrons.u.spline_lift(x_bnd_arr, y_check, z_check, squeeze_out=True)[0] + print(f"elec ux normal trace diff at {label}: max={xp.max(xp.abs(mms_vals - lift_vals)):.3e}") + + # comp=1: normal trace at y=0 and y=1 + for y_bnd, label in [(0.0, 'y=0'), (1.0, 'y=1')]: + y_bnd_arr = xp.array([y_bnd]) + mms_vals = mms_ion_u(x_check, y_bnd_arr, z_check)[1] + lift_vals = model.ions.u.spline_lift(x_check, y_bnd_arr, z_check, squeeze_out=True)[1] + print(f"ion uy normal trace diff at {label}: max={xp.max(xp.abs(mms_vals - lift_vals)):.3e}") + + mms_vals = mms_electron_u(x_check, y_bnd_arr, z_check)[1] + lift_vals = model.electrons.u.spline_lift(x_check, y_bnd_arr, z_check, squeeze_out=True)[1] + print(f"elec uy normal trace diff at {label}: max={xp.max(xp.abs(mms_vals - lift_vals)):.3e}") \ No newline at end of file diff --git a/examples/TwoFluidQuasiNeutralToy/R_Verification.py b/examples/TwoFluidQuasiNeutralToy/R_Verification.py new file mode 100644 index 000000000..f94380a24 --- /dev/null +++ b/examples/TwoFluidQuasiNeutralToy/R_Verification.py @@ -0,0 +1,217 @@ +import os +import glob +import cunumpy as xp +import matplotlib.pyplot as plt +from mpi4py import MPI + +from struphy.io.options import EnvironmentOptions, Time +from struphy.geometry import domains +from struphy.fields_background import equils +from struphy.topology import grids +from struphy.io.options import DerhamOptions +from struphy.initial.base import GenericPerturbation +from struphy import Simulation +from struphy.linear_algebra.solver import SolverParameters +from struphy.models.two_fluid_quasi_neutral_toy import TwoFluidQuasiNeutralToy + +# ------------------ parameters ------------------ +name = "runs/sim_restelli" + +R0 = 2.0 +a = 1.0 +ain = 0.1 +B0 = 10.0 +Bp = 12.5 +nu = 1.0 +nu_e = 0.01 +alpha = 0.1 +beta = 1.0 +eps = 1.0 +sigma = 1e-6 +dt = 1 +Tend = 1 +Nel = (8, 8, 1) +p = (1, 1, 1) +tol = 1e-5 + +env = EnvironmentOptions(sim_folder=name) +time_opts = Time(dt=dt, Tend=Tend) + +# ------------------ domain & equilibrium ------------------ +domain = domains.HollowTorus(a1=ain, a2=a, R0=R0) +equil = equils.CircularTokamak(a=a, R0=R0, B0=B0, Bp=Bp) +grid = grids.TensorProductGrid(num_elements=Nel) + +derham_opts = DerhamOptions( + degree=p, + bcs=(("dirichlet", "dirichlet"), None, None) +) + +# ------------------ manufactured solution ------------------ + +def _cylindrical(x, y, z): + R = xp.sqrt(x**2 + y**2) + R = xp.where(R == 0.0, 1e-9, R) + phi = xp.arctan2(-y, x) + Z = z + return R, phi, Z + +def mms_u_cartesian(x, y, z): + R, phi, Z = _cylindrical(x, y, z) + u_R = alpha/R * (-Z) / (a*R0/R) + beta * Bp/B0 * R0/(a*R) * Z + u_Z = alpha/R * (R - R0) / (a*R0/R) + beta * Bp/B0 * R0/(a*R) * (-(R - R0)) + u_phi = beta * Bp/B0 * R0/(a*R) * B0*Bp*a / (Bp*R0) + + u_R = alpha * R / (a*R0) * (-Z) + beta * Bp/B0 * R0/(a*R) * Z + u_Z = alpha * R / (a*R0) * (R - R0) + beta * Bp/B0 * R0/(a*R) * (-(R - R0)) + u_phi = beta * Bp/B0 * R0/(a*R) * (B0/Bp * a) + + # Transform to Cartesian via eq (5.14) + ux = xp.cos(phi) * u_R - R * xp.sin(phi) * u_phi + uy = -xp.sin(phi) * u_R - R * xp.cos(phi) * u_phi + uz = u_Z + return ux, uy, uz + +def mms_phi_cartesian(x, y, z): + R, phi, Z = _cylindrical(x, y, z) + # eq (5.22): phi_hat = 0.5 * a * B0 * alpha * ((R-R0)^2 + Z^2)/a^2 - 2/3) + phi_val = 0.5 * a * B0 * alpha * (((R - R0)**2 + Z**2) / a**2 - 2.0/3.0) + return phi_val + +# ------------------ source terms ------------------ + +def _omega_cartesian(x, y, z): + R, phi, Z = _cylindrical(x, y, z) + omega_Z = alpha * (R0 - 4*R) / (a*R0*R) - beta * Bp/B0 * R0**2 / (a*R**3) + ox = xp.zeros_like(x) + oy = xp.zeros_like(x) + oz = omega_Z + return ox, oy, oz + +def source_function_u(x, y, z): + ox, oy, oz = _omega_cartesian(x, y, z) + return nu * ox, nu * oy, nu * oz + +def source_function_ue(x, y, z): + ox, oy, oz = _omega_cartesian(x, y, z) + return nu_e * ox, nu_e * oy, nu_e * oz + +# ------------------ lifting (inhomogeneous Dirichlet on radial boundary) ------------------ +lifting_u = [ + GenericPerturbation(lambda x, y, z: mms_u_cartesian(x, y, z)[0], comp=0, given_in_basis="physical"), + GenericPerturbation(lambda x, y, z: mms_u_cartesian(x, y, z)[1], comp=1, given_in_basis="physical"), + GenericPerturbation(lambda x, y, z: mms_u_cartesian(x, y, z)[2], comp=2, given_in_basis="physical"), +] + +lifting_ue = [ + GenericPerturbation(lambda x, y, z: mms_u_cartesian(x, y, z)[0], comp=0, given_in_basis="physical"), + GenericPerturbation(lambda x, y, z: mms_u_cartesian(x, y, z)[1], comp=1, given_in_basis="physical"), + GenericPerturbation(lambda x, y, z: mms_u_cartesian(x, y, z)[2], comp=2, given_in_basis="physical"), +] + +# ------------------ model ------------------ +model = TwoFluidQuasiNeutralToy() + +model.propagators.qn_full.options = model.propagators.qn_full.Options( + nu=nu, + nu_e=nu_e, + eps_norm=eps, + stab_sigma=sigma, + source_u=source_function_u, + source_ue=source_function_ue, + solver='gmres', + solver_params=SolverParameters(verbose=True, info=True, tol=tol), +) + +model.ions.u.lifting_function = lifting_u +model.electrons.u.lifting_function = lifting_ue + +# ------------------ simulation ------------------ +sim = Simulation( + model=model, + params_path=__file__, + env=env, + time_opts=time_opts, + domain=domain, + equil=equil, + grid=grid, + derham_opts=derham_opts, +) + +# ------------------ run ------------------ +if __name__ == "__main__": + # sim.run(verbose=True) + + if MPI.COMM_WORLD.Get_rank() == 0: + sim.pproc(verbose=True) + sim.load_plotting_data(verbose=True) + + simdata = sim.plotting_data + os.makedirs(f'{name}/plots', exist_ok=True) + for f in glob.glob(f'{name}/plots/*.png'): + os.remove(f) + + e1 = xp.linspace(0, 1, 40) + e2 = xp.linspace(0, 1, 40) + e3 = xp.array([0.5]) + x_phys, y_phys, z_phys = domain(e1, e2, e3, squeeze_out=True) + + theta_bnd = xp.linspace(0, 1, 200) + x_inner, y_inner, _ = domain(xp.array([0.0]), theta_bnd, xp.array([0.5]), squeeze_out=True) + x_outer, y_outer, _ = domain(xp.array([1.0]), theta_bnd, xp.array([0.5]), squeeze_out=True) + + def _add_domain_boundary(ax): + ax.plot(x_inner, y_inner, 'w-', linewidth=0.8) + ax.plot(x_outer, y_outer, 'w-', linewidth=0.8) + + prop = model.propagators.qn_full + + for label, spline, src_fn, comp in [ + ('ion_source_x', prop._src_u, prop.options.source_u, 0), + ('ion_source_y', prop._src_u, prop.options.source_u, 1), + ('ion_source_z', prop._src_u, prop.options.source_u, 2), + ('electron_source_x', prop._src_ue, prop.options.source_ue, 0), + ('electron_source_y', prop._src_ue, prop.options.source_ue, 1), + ('electron_source_z', prop._src_ue, prop.options.source_ue, 2), + ]: + if spline is None: + print(f" {label}: None, skipping") + continue + + vals_proj = spline(e1, e2, e3, squeeze_out=True)[comp] + vals_ref = src_fn(x_phys, y_phys, z_phys)[comp] + + fig, axes = plt.subplots(1, 2, figsize=(10, 4)) + im0 = axes[0].contourf(x_phys, y_phys, vals_proj, levels=50); axes[0].set_title('projected (FE)'); plt.colorbar(im0, ax=axes[0]); _add_domain_boundary(axes[0]) + im1 = axes[1].contourf(x_phys, y_phys, vals_ref, levels=50); axes[1].set_title('reference (analytical)'); plt.colorbar(im1, ax=axes[1]); _add_domain_boundary(axes[1]) + fig.suptitle(label) + out = f'{name}/plots/source_{label}.png' + plt.savefig(out, dpi=300) + plt.close(fig) + print(f" -> saved {out}") + + for t in simdata.spline_values.ions.u_log.data.keys(): + u_ions = simdata.spline_values.ions.u_log.data[t] + u_electrons = simdata.spline_values.electrons.u_log.data[t] + phi_num = simdata.spline_values.em_fields.phi_log.data[t] + + mms_ux, mms_uy, mms_uz = mms_u_cartesian(x_phys, y_phys, z_phys) + mms_phi_val = mms_phi_cartesian(x_phys, y_phys, z_phys) + + for num, mms, lbl in [ + (u_ions[0][:, :, 0], mms_ux, 'u_ix'), + (u_ions[1][:, :, 0], mms_uy, 'u_iy'), + (u_ions[2][:, :, 0], mms_uz, 'u_iz'), + (u_electrons[0][:, :, 0], mms_ux, 'u_ex'), + (u_electrons[1][:, :, 0], mms_uy, 'u_ey'), + (u_electrons[2][:, :, 0], mms_uz, 'u_ez'), + (phi_num[0][:, :, 0], mms_phi_val, 'phi'), + ]: + fig, axes = plt.subplots(1, 3, figsize=(15, 4)) + im0 = axes[0].contourf(x_phys, y_phys, num, levels=50); axes[0].set_title('numerical'); plt.colorbar(im0, ax=axes[0]); _add_domain_boundary(axes[0]) + im1 = axes[1].contourf(x_phys, y_phys, mms, levels=50); axes[1].set_title('MMS'); plt.colorbar(im1, ax=axes[1]); _add_domain_boundary(axes[1]) + im2 = axes[2].contourf(x_phys, y_phys, num - mms, levels=50); axes[2].set_title('difference'); plt.colorbar(im2, ax=axes[2]); _add_domain_boundary(axes[2]) + fig.suptitle(f'{lbl} at t={t:.4f}') + out = f'{name}/plots/{lbl}_{t:.4f}.png' + plt.savefig(out, dpi=300) + plt.close(fig) \ No newline at end of file diff --git a/feectools b/feectools index 8c88dec79..ae6859fcb 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 8c88dec79510b315024d4b7e0ccc28e76ad8c9e7 +Subproject commit ae6859fcb16c765f7bb7cabb82eb6268d1967014 diff --git a/src/struphy/feec/linear_operators.py b/src/struphy/feec/linear_operators.py index 940bcd4d5..1cb513ca1 100644 --- a/src/struphy/feec/linear_operators.py +++ b/src/struphy/feec/linear_operators.py @@ -304,21 +304,31 @@ class BoundaryOperator(LinOpWithTransp): Parameters ---------- vector_space : feectools.linalg.basic.VectorSpace - The vector space associated to the operator. + The vector space of the domain (input). space_id : str Symbolic space ID of vector_space (H1, Hcurl, Hdiv, L2 or H1vec). dirichlet_bc : tuple[tuple[bool]] Whether to apply homogeneous Dirichlet boundary conditions (at left or right boundary in each direction). + + codomain : feectools.linalg.basic.VectorSpace, optional + The vector space of the codomain (output). If given, the operator maps between two different spaces + (e.g. unconstrained to constrained). If None, domain and codomain are the same. """ - def __init__(self, vector_space, space_id, dirichlet_bc): + def __init__(self, vector_space, space_id, dirichlet_bc, codomain=None): assert isinstance(vector_space, VectorSpace) assert isinstance(space_id, str) self._domain = vector_space - self._codomain = vector_space + if codomain is not None: + assert isinstance(codomain, VectorSpace) + self._codomain = codomain + self._cross_space = True + else: + self._codomain = vector_space + self._cross_space = False self._dtype = vector_space.dtype self._space_id = space_id @@ -491,14 +501,16 @@ def dot(self, v, out=None): assert isinstance(v, Vector) assert v.space == self._domain - if out is None: - out = v.copy() - else: + if out is not None: assert isinstance(out, Vector) assert out.space == self._codomain v.copy(out=out) + elif self._cross_space: + out = self._codomain.zeros() + v.copy(out=out) + else: + out = v.copy() - # apply boundary conditions to output vector apply_essential_bc_to_array(self._space_id, out, self.bc) return out @@ -507,4 +519,4 @@ def transpose(self, conjugate=False): """ Returns the transposed operator. """ - return BoundaryOperator(self._domain, self._space_id, self.bc) + return BoundaryOperator(self._codomain, self._space_id, self.bc, codomain=self._domain) diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 6b722b081..db16037d3 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -261,6 +261,13 @@ def boundary_spline(self) -> SplineFunction | None: self._boundary_spline = None return self._boundary_spline + @property + def spline_full(self) -> SplineFunction | None: + """Full solution spline (lifting + zero-BC part) in the unconstrained space. Only allocated if lifting_function is not None.""" + if not hasattr(self, "_spline_full"): + self._spline_full = None + return self._spline_full + @property def boundary_op(self) -> BoundaryOperator | None: """The boundary operator, used for the lifting of boundary conditions. Only allocated if lifting_function is not None.""" @@ -268,6 +275,14 @@ def boundary_op(self) -> BoundaryOperator | None: self._boundary_op = None return self._boundary_op + @property + def boundary_op_lift(self) -> BoundaryOperator | None: + """Boundary operator mapping from the unconstrained (lifted) space to the constrained space. + Only allocated if lifting_function is not None.""" + if not hasattr(self, "_boundary_op_lift"): + self._boundary_op_lift = None + return self._boundary_op_lift + @property def derham_lift(self) -> Derham | None: """The Derham object for the lifting function. Only allocated if lifting_function is not None.""" @@ -324,23 +339,36 @@ def allocate( f"Lifting of boundary conditions can only be applied if at least one homogenous Dirichlet boundary condition is present in the Derham object, but here {derham.bcs = }" ) - # create another Derham object with the same options but with homogenous Dirichlet BCs replaced by free BCs, to be used for the lifting function + # normalise to list + lifting_list = self.lifting_function if isinstance(self.lifting_function, list) else [self.lifting_function] + + # validation + if self.space in {"H1", "L2"}: + if len(lifting_list) > 1: + raise ValueError("H1/L2 lifting only accepts a single Perturbation, not a list.") + elif self.space in {"Hcurl", "Hdiv", "H1vec"}: + if len(lifting_list) > 3: + raise ValueError("Hdiv/Hcurl/H1vec lifting accepts at most 3 Perturbations (one per component).") + comps = [ptb.comp for ptb in lifting_list] + if len(comps) != len(set(comps)): + raise ValueError(f"Each component may only appear once in the lifting list, got {comps}.") + + # create unconstrained Derham dct = derham.to_dict() bcs_lift = list(dct["options"]["bcs"]) for i, bc in enumerate(bcs_lift): if bc is not None: - bcn = list(bc) # convert tuple to list to allow modification + bcn = list(bc) if bcn[0] == "dirichlet": bcn[0] = "free" if bcn[1] == "dirichlet": bcn[1] = "free" - bcn = tuple(bcn) # convert back to tuple - bcs_lift[i] = bcn - dct["options"]["bcs"] = tuple(bcs_lift) # convert list back to tuple + bcs_lift[i] = tuple(bcn) + dct["options"]["bcs"] = tuple(bcs_lift) self._derham_lift = Derham.from_dict(dct, comm=derham.comm) - # spline function for the lifting function + # spline function for the lifting self._spline_lift = self.derham_lift.create_spline_function( name=self.__name__ + "_lift" if self.__name__ is not None else None, space_id=self.space, @@ -349,50 +377,60 @@ def allocate( verbose=verbose, ) - # project lifting function to spline space - ptb = self.lifting_function + # spline function for unconstrained solution + self._spline_full = self.derham_lift.create_spline_function( + name=self.__name__ + "_full" if self.__name__ is not None else None, + space_id=self.space, + domain=domain, + equil=equil, + verbose=verbose, + ) - if self.space in { - "H1", - "L2", - }: # TODO: this is a copy-paste from SplineFunction.initialize_coeffs(), to be unified + # project each perturbation and accumulate into spline_lift + if self.space in {"H1", "L2"}: + ptb = lifting_list[0] if ptb.given_in_basis is None: ptb.given_in_basis = "0" - fun = TransformedPformComponent( ptb, ptb.given_in_basis, derham.space_to_form[self.space], domain=domain, ) + self.spline_lift.vector += self.derham_lift.projectors[derham.space_to_form[self.space]](fun) + elif self.space in {"Hcurl", "Hdiv", "H1vec"}: fun_vec = [None] * 3 - fun_vec[ptb.comp] = ptb - - if ptb.given_in_basis is None: - ptb.given_in_basis = "v" - # pullback callable for each component - fun = [] - for comp in range(3): - fun += [ - TransformedPformComponent( - fun_vec, - ptb.given_in_basis, - derham.space_to_form[self.space], - comp=comp, - domain=domain, - ), - ] - - # peform projection - self.spline_lift.vector += self.derham_lift.projectors[derham.space_to_form[self.space]](fun) - - # other helper objects for the lifting of boundary conditions + for ptb in lifting_list: + if fun_vec[ptb.comp] is not None: + raise ValueError(f"Component {ptb.comp} assigned more than once in lifting list.") + fun_vec[ptb.comp] = ptb + if ptb.given_in_basis is None: + ptb.given_in_basis = "v" + + fun = [ + TransformedPformComponent( + fun_vec, + fun_vec[comp].given_in_basis if fun_vec[comp] is not None else lifting_list[0].given_in_basis, + derham.space_to_form[self.space], + comp=comp, + domain=domain, + ) + for comp in range(3) + ] + self.spline_lift.vector += self.derham_lift.projectors[derham.space_to_form[self.space]](fun) + + # other helper objects self._spline_0 = self.spline_lift.copy() - self.spline_0.vector[:] = self.spline_lift.vector[:] + self.spline_lift.vector.copy(out=self.spline_0.vector) self._boundary_spline = self.spline_lift.copy() + self._boundary_op = BoundaryOperator(self.spline_lift.space, self.space, derham.dirichlet_bc) + self._boundary_op_lift = BoundaryOperator( + self.spline_lift.space, self.space, derham.dirichlet_bc, codomain=self._spline.space + ) + self.compute_boundary_spline() def compute_boundary_spline(self, spline_lift: SplineFunction | None = None): @@ -405,7 +443,7 @@ def compute_boundary_spline(self, spline_lift: SplineFunction | None = None): # set new boundary spline diff_vec = spline_lift.vector - self.spline_0.vector - self.boundary_spline.vector[:] = diff_vec[:] + diff_vec.copy(out=self.boundary_spline.vector) class PICVariable(Variable): diff --git a/src/struphy/physics/physics.py b/src/struphy/physics/physics.py index 3ffff52b1..2bb40a861 100644 --- a/src/struphy/physics/physics.py +++ b/src/struphy/physics/physics.py @@ -85,6 +85,13 @@ def j(self): raise AttributeError("Must call Units.derive_units() to get full set of units.") return self._j + @property + def nu(self): + """Unit of dynamic viscosity in kg/(m·s).""" + if not hasattr(self, "_nu"): + raise AttributeError("Must call Units.derive_units() to get full set of units.") + return self._nu + def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk: int = None, verbose=False): """Derive the remaining units from the base units, velocity scale and bulk species' A and Z.""" @@ -129,6 +136,9 @@ def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk # current density (A/m^2) self._j = con.e * self.n * self.v + # dynamic viscosity (kg/(m·s)) + self._nu = A_bulk * con.mH * self.n * self.x * self.v if A_bulk is not None else None + # print to screen if verbose and MPI.COMM_WORLD.Get_rank() == 0: units_used = ( @@ -141,6 +151,7 @@ def derive_units(self, velocity_scale: str = "light", A_bulk: int = None, Z_bulk " bar", " kg/m³", " A/m²", + " kg/(m·s)", ) logger.info("") for (k, v), u in zip(self.__dict__.items(), units_used): diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 3ea76df93..062b72737 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -118,6 +118,12 @@ def update_feec_variables(self, **new_coeffs): old = old_var.spline.vector assert new.space == old.space + # update full solution spline (lifting + zero-BC part) if present + if old_var.spline_full is not None: + new.copy(out=old_var.spline_full.vector) + if old_var.boundary_spline is not None: + old_var.spline_full.vector += old_var.boundary_spline.vector + # calculate maximum of difference abs(new - old) diffs[var] = xp.max(xp.abs(new.toarray() - old.toarray())) diff --git a/src/struphy/propagators/two_fluid_quasi_neutral_full.py b/src/struphy/propagators/two_fluid_quasi_neutral_full.py index 741a0d58e..5c55ef13c 100644 --- a/src/struphy/propagators/two_fluid_quasi_neutral_full.py +++ b/src/struphy/propagators/two_fluid_quasi_neutral_full.py @@ -11,6 +11,7 @@ from struphy.feec.basis_projection_ops import BasisProjectionOperators from struphy.feec.mass import L2Projector, WeightedMassOperators +from struphy.geometry.utilities import TransformedPformComponent from struphy.io.options import LiteralOptions from struphy.linear_algebra.solver import SolverParameters from struphy.models.variables import FEECVariable @@ -106,10 +107,6 @@ class Options: Electron viscosity coefficient. eps_norm : float, default=1e-3 Normalization/scaling parameter in Lorentz coupling terms. - boundary_data_u : dict[tuple[int, int], Callable] or None, default=None - Inhomogeneous Dirichlet data for ion velocity faces. - boundary_data_ue : dict[tuple[int, int], Callable] or None, default=None - Inhomogeneous Dirichlet data for electron velocity faces. source_u : Callable or None, default=None Source term for ion momentum equation. source_ue : Callable or None, default=None @@ -124,34 +121,32 @@ class Options: nu: float = 1.0 nu_e: float = 1.0 - eps_norm: float = 1e-3 - - boundary_data_u: dict[tuple[int, int], Callable] | None = None - boundary_data_ue: dict[tuple[int, int], Callable] | None = None + eps_norm: float | None = None source_u: Callable | None = None source_ue: Callable | None = None - stab_sigma: float | None = None - + stab_sigma: float = 0.0 solver: LiteralOptions.OptsGenSolver = "gmres" solver_params: SolverParameters | None = None def __post_init__(self): + # --- warn if no source terms --- + if self.source_u is None: + warn("No source_u specified — defaulting to zero.") + if self.source_ue is None: + warn("No source_ue specified — defaulting to zero.") + if self.eps_norm is None: + warn("No eps_norm specified — will default to ion cyclotron parameter epsilon in allocate.") + # --- physical parameter sanity checks --- if self.nu < 0: raise ValueError(f"nu must be non-negative, got {self.nu}") if self.nu_e < 0: raise ValueError(f"nu_e must be non-negative, got {self.nu_e}") - if self.eps_norm <= 0: + if self.eps_norm is not None and self.eps_norm <= 0: raise ValueError(f"eps_norm must be positive, got {self.eps_norm}") - # --- warn if no source terms --- - if self.source_u is None: - warn("No source_u specified — defaulting to zero.") - if self.source_ue is None: - warn("No source_ue specified — defaulting to zero.") - # --- defaults --- if self.stab_sigma is None: warn("stab_sigma not specified, defaulting to 0.0") @@ -163,7 +158,8 @@ def __post_init__(self): @property def options(self) -> Options: - assert hasattr(self, "_options"), "Options not set." + if not hasattr(self, "_options"): + self._options = self.Options() return self._options @options.setter @@ -175,45 +171,6 @@ def options(self, new): logger.info(f" {k}: {v}") self._options = new - # ========================================================================= - ### Boundary condition helpers - # ========================================================================= - - def _get_dirichlet_faces(self): - """Infer which faces have Dirichlet BCs by comparing derham and derham_v0. - - A face is Dirichlet if it is unclamped in derham but clamped in derham_v0 - (i.e. lifting is True there). - """ - faces = [] - derham = self.derham - derham_v0 = derham - - if derham_v0 is None: - return faces - - bc = derham.dirichlet_bc - bc_v0 = derham_v0.dirichlet_bc - - for d in range(3): - if derham.spl_kind[d]: - continue # periodic axis, no Dirichlet - for s, side in enumerate((-1, 1)): - # clamped in v0 but not in derham => this is a lifted (inhom Dirichlet) face - unclamped = not bc[d][s] - clamped_v0 = bc_v0[d][s] if bc_v0 is not None else False - if unclamped and clamped_v0: - faces.append((d, side)) - # clamped in both => homogeneous Dirichlet, also need to zero DOFs - elif bc[d][s] and clamped_v0: - faces.append((d, side)) - return faces - - def _apply_essential_bc(self, vec): - """Zero out Dirichlet DOFs, inferred from derham vs derham_v0.""" - for d, side in self._dirichlet_faces: - apply_essential_bc_stencil(vec[0], axis=d, ext=side, order=0) - # ========================================================================= ### Allocate # ========================================================================= @@ -224,95 +181,175 @@ def allocate(self, verbose=False): self._rank = self.derham.comm.Get_rank() if self.derham.comm is not None else 0 self._dt = None - # ---- v0 de Rham complex (from derham.derham_v0) ---------------------- - self._derham_v0 = self.derham + if self.options.eps_norm is None: + self._options.eps_norm = self.variables.u.species.equation_params.epsilon + + # ---- lifting (derham_lift is unconstrained, self.derham is constrained) --- + self._has_lifting_u = self.variables.u.derham_lift is not None + self._has_lifting_ue = self.variables.ue.derham_lift is not None + + self._derham_lift_u = self.variables.u.derham_lift if self._has_lifting_u else self.derham + self._derham_lift_ue = self.variables.ue.derham_lift if self._has_lifting_ue else self.derham + + # ---- solution splines (constrained) and u in unconstrained space ----- + self._u_0 = self.derham.create_spline_function("u", space_id="Hdiv") - self._mass_ops_v0 = WeightedMassOperators( - self._derham_v0, + # boundary splines (u', ue') in unconstrained space — zero vectors if no lifting + self._boundary_spline_u = ( + self.variables.u.boundary_spline.vector + if self._has_lifting_u + else self._derham_lift_u.coeff_spaces["2"].zeros() + ) + self._boundary_spline_ue = ( + self.variables.ue.boundary_spline.vector + if self._has_lifting_ue + else self._derham_lift_ue.coeff_spaces["2"].zeros() + ) + + # boundary operators + self._b_op_u = ( + self.variables.u.boundary_op_lift + if self._has_lifting_u + else IdentityOperator(self.derham.coeff_spaces["2"]) + ) + self._b_op_ue = ( + self.variables.ue.boundary_op_lift + if self._has_lifting_ue + else IdentityOperator(self.derham.coeff_spaces["2"]) + ) + + # pre-allocated RHS vectors (constrained, after boundary operator) + self._rhs_vec_u = self.derham.create_spline_function("rhs_vec_u", space_id="Hdiv") + self._rhs_vec_ue = self.derham.create_spline_function("rhs_vec_ue", space_id="Hdiv") + self._rhs_vec_phi = self.derham.create_spline_function("rhs_vec_phi", space_id="L2") + + self._div_boundary_u = self.derham.create_spline_function("div_boundary_u", space_id="L2") + self._div_boundary_ue = self.derham.create_spline_function("div_boundary_ue", space_id="L2") + + # ---- source terms projected onto unconstrained space ----------------- + self._src_u = self._derham_lift_u.create_spline_function("rhs_u", space_id="Hdiv") + self._src_ue = self._derham_lift_ue.create_spline_function("rhs_ue", space_id="Hdiv") + + for rhs, source, derham_lift in [ + (self._src_u, self.options.source_u, self._derham_lift_u), + (self._src_ue, self.options.source_ue, self._derham_lift_ue), + ]: + if source is not None: + fun_vec = [lambda x, y, z, f=source, c=c: f(x, y, z)[c] for c in range(3)] + fun = [ + TransformedPformComponent( + fun_vec, + "physical", + "2", + comp=comp, + domain=self.domain, + ) + for comp in range(3) + ] + rhs.vector = derham_lift.projectors["2"](fun) + + # ---- unconstrained mass/basis operators (for RHS assembly) ----------- + + self._mass_ops_lift_u = WeightedMassOperators( + self._derham_lift_u, + self.domain, + eq_mhd=self.mass_ops.eq_mhd, + ) + self._mass_ops_lift_ue = WeightedMassOperators( + self._derham_lift_ue, self.domain, eq_mhd=self.mass_ops.eq_mhd, ) - self._basis_ops_v0 = BasisProjectionOperators( - self._derham_v0, + self._basis_ops_lift_u = BasisProjectionOperators( + self._derham_lift_u, self.domain, verbose=self.options.solver_params.verbose, eq_mhd=self.basis_ops.weights["eq_mhd"], ) + self._basis_ops_lift_ue = BasisProjectionOperators( + self._derham_lift_ue, + self.domain, + verbose=self.options.solver_params.verbose, + eq_mhd=self.basis_ops.weights["eq_mhd"], + ) + + self._M2_u = self._mass_ops_lift_u.M2 + self._M2B_u = -self._mass_ops_lift_u.M2B + self._div_u = self._derham_lift_u.div + self._curl_u = self._derham_lift_u.curl + self._S21_u = self._basis_ops_lift_u.S21 + + self._lapl_u = ( + self._div_u.T @ self._mass_ops_lift_u.M3 @ self._div_u + + self._S21_u.T @ self._curl_u.T @ self._M2_u @ self._curl_u @ self._S21_u + ) + + self._A11_u = -self._M2B_u / self.options.eps_norm + self.options.nu * self._lapl_u - # ---- Dirichlet faces (inferred from derham vs derham_v0) ------------- + self._M2_ue = self._mass_ops_lift_ue.M2 + self._M2B_ue = -self._mass_ops_lift_ue.M2B + self._div_ue = self._derham_lift_ue.div + self._curl_ue = self._derham_lift_ue.curl + self._S21_ue = self._basis_ops_lift_ue.S21 - self._dirichlet_faces = self._get_dirichlet_faces() + self._lapl_ue = ( + self._div_ue.T @ self._mass_ops_lift_ue.M3 @ self._div_ue + + self._S21_ue.T @ self._curl_ue.T @ self._M2_ue @ self._curl_ue @ self._S21_ue + ) + + self._A22_ue = ( + self.options.stab_sigma * IdentityOperator(self._derham_lift_ue.coeff_spaces["2"]) + + self._M2B_ue / self.options.eps_norm + + self.options.nu_e * self._lapl_ue + ) - # ---- unconstrained operators (for RHS assembly) ---------------------- + # ---- constrained operators (for system matrix, built from self.derham) --- self._M2 = self.mass_ops.M2 + self._M3 = self.mass_ops.M3 self._M2B = -self.mass_ops.M2B self._div = self.derham.div self._curl = self.derham.curl self._S21 = self.basis_ops.S21 - self._lapl = ( - self._div.T @ self.mass_ops.M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21 + self._lapl_v0 = ( + self._div.T @ self._M3 @ self._div + self._S21.T @ self._curl.T @ self._M2 @ self._curl @ self._S21 ) - self._A11 = -self._M2B / self.options.eps_norm + self.options.nu * self._lapl + self._A11 = -self._M2B / self.options.eps_norm + self.options.nu * self._lapl_v0 self._A22 = ( - -self.options.stab_sigma * IdentityOperator(self.derham.V2) + self.options.stab_sigma * IdentityOperator(self.derham.coeff_spaces["2"]) + self._M2B / self.options.eps_norm - + self.options.nu_e * self._lapl - ) - - # ---- constrained operators (for system matrix) ----------------------- - - self._M2_v0 = self._mass_ops_v0.M2 - self._M3_v0 = self._mass_ops_v0.M3 - self._M2B_v0 = -self._mass_ops_v0.M2B - self._div_v0 = self._derham_v0.div - self._curl_v0 = self._derham_v0.curl - self._S21_v0 = self._basis_ops_v0.S21 - - self._lapl_v0 = ( - self._div_v0.T @ self._M3_v0 @ self._div_v0 - + self._S21_v0.T @ self._curl_v0.T @ self._M2_v0 @ self._curl_v0 @ self._S21_v0 - ) - - self._A11_v0 = -self._M2B_v0 / self.options.eps_norm + self.options.nu * self._lapl_v0 - self._A22_v0 = ( - -self.options.stab_sigma * IdentityOperator(self._derham_v0.V2) - + self._M2B_v0 / self.options.eps_norm + self.options.nu_e * self._lapl_v0 ) # ---- block saddle-point system ---------------------------------------- - self._block_domain_v0 = BlockVectorSpace(self._derham_v0.V2, self._derham_v0.V2) - self._block_codomain_v0 = self._block_domain_v0 - self._block_codomain_B_v0 = self._derham_v0.V3 + self._block_domain = BlockVectorSpace(self.derham.coeff_spaces["2"], self.derham.coeff_spaces["2"]) + self._block_codomain_B = self.derham.coeff_spaces["3"] - self._B1_v0 = -self._M3_v0 @ self._div_v0 - self._B2_v0 = self._M3_v0 @ self._div_v0 + self._B1 = -self._M3 @ self._div + self._B2 = self._M3 @ self._div - self._B_v0 = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_B_v0, blocks=[[self._B1_v0, self._B2_v0]] - ) + self._B = BlockLinearOperator(self._block_domain, self._block_codomain_B, blocks=[[self._B1, self._B2]]) - self._block_domain_M = BlockVectorSpace(self._block_domain_v0, self._block_codomain_B_v0) + self._block_domain_M = BlockVectorSpace(self._block_domain, self._block_codomain_B) _A_init = BlockLinearOperator( - self._block_domain_v0, self._block_codomain_v0, blocks=[[self._A11_v0, None], [None, self._A22_v0]] + self._block_domain, self._block_domain, blocks=[[self._A11, None], [None, self._A22]] ) _M_init = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, blocks=[[_A_init, self._B_v0.T], [self._B_v0, None]] + self._block_domain_M, self._block_domain_M, blocks=[[_A_init, self._B.T], [self._B, None]] ) if self.options.solver in get_args(LiteralOptions.OptsSaddlePointSolver): self._Minv = inverse( _M_init, self.options.solver, - A11=self._A11_v0, - A22=self._A22_v0, - B1=self._B1_v0, - B2=self._B2_v0, + A11=self._A11, + A22=self._A22, + B1=self._B1, + B2=self._B2, recycle=self.options.solver_params.recycle, tol=self.options.solver_params.tol, maxiter=self.options.solver_params.maxiter, @@ -328,140 +365,63 @@ def allocate(self, verbose=False): verbose=self.options.solver_params.verbose, ) - # ---- projector ------------------------------------------------------- - - self._projector = L2Projector(space_id="Hdiv", mass_ops=self.mass_ops) - - # ---- solution spline functions (unconstrained) ----------------------- - - self._u = self.derham.create_spline_function("u", space_id="Hdiv") - self._ue = self.derham.create_spline_function("ue", space_id="Hdiv") - self._phi = self.derham.create_spline_function("phi", space_id="L2") - - # ---- BC lifts (unconstrained) ---------------------------------------- - - self._u_prime = self.derham.create_spline_function("u_prime", space_id="Hdiv") - self._ue_prime = self.derham.create_spline_function("ue_prime", space_id="Hdiv") - - for u_prime, boundary_data in [ - (self._u_prime, self.options.boundary_data_u), - (self._ue_prime, self.options.boundary_data_ue), - ]: - if boundary_data is None: - continue - for (d, side), f_bc in boundary_data.items(): - if (d, side) in self._dirichlet_faces: - bc_pulled = lambda *etas, f=f_bc: self.domain.pull( - [ - lambda x, y, z, f=f: f(x, y, z)[0], - lambda x, y, z, f=f: f(x, y, z)[1], - lambda x, y, z, f=f: f(x, y, z)[2], - ], - *etas, - kind="2", - ) - _vec = self._projector( - [ - lambda *etas: bc_pulled(*etas)[0], - lambda *etas: bc_pulled(*etas)[1], - lambda *etas: bc_pulled(*etas)[2], - ] - ) - for d2, side2 in self._dirichlet_faces: - if (d2, side2) != (d, side): - apply_essential_bc_stencil(_vec[0], axis=d2, ext=side2, order=0) - u_prime.vector += _vec - - self._u_prime_v0 = self._derham_v0.create_spline_function("u_prime_v0", space_id="Hdiv") - self._ue_prime_v0 = self._derham_v0.create_spline_function("ue_prime_v0", space_id="Hdiv") - - self._u_prime_v0.vector = self._u_prime.vector - self._ue_prime_v0.vector = self._ue_prime.vector - - # ---- projected source terms (unconstrained) -------------------------- - - self._rhs_u = self.derham.create_spline_function("rhs_u", space_id="Hdiv") - self._rhs_ue = self.derham.create_spline_function("rhs_ue", space_id="Hdiv") - - for rhs, source in [(self._rhs_u, self.options.source_u), (self._rhs_ue, self.options.source_ue)]: - if source is not None: - src_pulled = lambda *etas, f=source: self.domain.pull( - [ - lambda x, y, z, f=f: f(x, y, z)[0], - lambda x, y, z, f=f: f(x, y, z)[1], - lambda x, y, z, f=f: f(x, y, z)[2], - ], - *etas, - kind="2", - ) - rhs.vector = self._projector.get_dofs( - [ - lambda *etas: src_pulled(*etas)[0], - lambda *etas: src_pulled(*etas)[1], - lambda *etas: src_pulled(*etas)[2], - ] - ) - - # ---- pre-allocated RHS vectors (v0, reused each time step) ----------- - - self._rhs_vec_u = self._derham_v0.create_spline_function("rhs_vec_u", space_id="Hdiv") - self._rhs_vec_ue = self._derham_v0.create_spline_function("rhs_vec_ue", space_id="Hdiv") + self._RHS = BlockVector( + self._block_domain_M, + blocks=[ + BlockVector(self._block_domain, blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]), + self._rhs_vec_phi.vector, + ], + ) + self._SOL = self._block_domain_M.zeros() # ========================================================================= ### Time step # ========================================================================= - def __call__(self, dt): - # --- copy current state --- - self._u.vector = self.variables.u.spline.vector - self._ue.vector = self.variables.ue.spline.vector - - # --- rebuild system matrix if dt changed --- - if dt != self._dt: # TODO change uzawa A11 block too - self._dt = dt - _A = BlockLinearOperator( - self._block_domain_v0, - self._block_codomain_v0, - blocks=[[self._A11_v0 + self._M2_v0 / dt, None], [None, self._A22_v0]], - ) + # --- copy current homogeneous solution --- + self._u_0.vector = self.variables.u.spline.vector - _M = BlockLinearOperator( - self._block_domain_M, self._block_domain_M, blocks=[[_A, self._B_v0.T], [self._B_v0, None]] + # --- assemble RHS fully in unconstrained space, then enforce essential BCs --- + self._rhs_vec_u.vector = ( + self._b_op_u.dot( + ( + self._M2_u.dot(self._src_u.vector) + - self._A11_u.dot(self._boundary_spline_u) + - self._M2_u.dot(self._boundary_spline_u) / dt + ) ) - self._Minv.linop = _M + + self._M2.dot(self._u_0.vector) / dt + ) - # --- assemble RHS in unconstrained space, then zero boundary DOFs --- - # ion: F1 = rhs_u + M2/dt * u - (A11 + M2/dt) * u' - # electron: F2 = rhs_ue - A22 * ue' - self._rhs_vec_u.vector = ( - self._rhs_u.vector # TODO boundary operator - + self._M2.dot(self._u.vector) / dt - - self._A11.dot(self._u_prime.vector) - - self._M2.dot(self._u_prime.vector) / dt + self._rhs_vec_ue.vector = self._b_op_ue.dot( + (self._M2_ue.dot(self._src_ue.vector) - self._A22_ue.dot(self._boundary_spline_ue)) ) - self._rhs_vec_ue.vector = self._rhs_ue.vector - self._A22.dot(self._ue_prime.vector) - self._apply_essential_bc(self._rhs_vec_u.vector) - self._apply_essential_bc(self._rhs_vec_ue.vector) + self._div_boundary_u.vector = self._div_u.dot(self._boundary_spline_u) + self._div_boundary_ue.vector = self._div_ue.dot(self._boundary_spline_ue) + + self._rhs_vec_phi.vector = self.mass_ops.M3.dot(self._div_boundary_u.vector) - self.mass_ops.M3.dot( + self._div_boundary_ue.vector + ) # --- build block RHS and solve --- - _F = BlockVector(self._block_domain_v0, blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]) - _RHS = BlockVector(self._block_domain_M, blocks=[_F, self._block_codomain_B_v0.zeros()]) + self._Minv.dot( + BlockVector( + self._block_domain_M, + blocks=[ + BlockVector(self._block_domain, blocks=[self._rhs_vec_u.vector, self._rhs_vec_ue.vector]), + self._rhs_vec_phi.vector, + ], + ), + out=self._SOL, + ) - _sol = self._Minv.dot(_RHS) info = self._Minv.get_info() - # --- reconstruct full solution: u = u_0 + u' --- - self._u.vector = _sol[0][0] + self._u_prime_v0.vector - self._ue.vector = _sol[0][1] + self._ue_prime_v0.vector - self._phi.vector = _sol[1] - # --- update FEEC variables --- - max_diffs = self.update_feec_variables(u=self._u.vector, ue=self._ue.vector, phi=self._phi.vector) + max_diffs = self.update_feec_variables(u=self._SOL[0][0], ue=self._SOL[0][1], phi=self._SOL[1]) if self.options.solver_params.info and self._rank == 0: - logger.info(f"Status: {info['success']}, Iterations: {info['niter']}") - logger.info(f"Max diffs: {max_diffs}") - logger.info(f"Status: {info['success']}, Iterations: {info['niter']}") - logger.info(f"Max diffs: {max_diffs}") + print(f"Status: {info['success']}, Iterations: {info['niter']}") + print(f"Max diffs: {max_diffs}") diff --git a/src/struphy/utils/utils.py b/src/struphy/utils/utils.py index 39ed325ef..502d662d0 100644 --- a/src/struphy/utils/utils.py +++ b/src/struphy/utils/utils.py @@ -109,7 +109,7 @@ def kernels_to_txt(kernels: list, output: str): # logger.info(f"kernels written to {output}.") -def check_option(opt, *options): +def check_option(opt: str | list[str], *options): """Check if opt is contained in options; if opt is a list, checks for each element.""" opts = [] for o in options: